分享一个openssh自动升级的python工具源代码

#!/usr/bin/env python3# -*- coding: utf-8 -*-"""OpenSSH 离线升级工具支持交互式输入路径或命令行参数指定兼容 Python 3.6+"""import osimport sysimport subprocessimport argparseimport tarfileimport shutilfrom pathlib import Path# 颜色输出class Colors:    RED = '\033[31m'    GREEN = '\033[32m'    YELLOW = '\033[33m'    BLUE = '\033[34m'    NC = '\033[0m'def print_colored(text, color=Colors.NC):    print(f"{color}{text}{Colors.NC}")def run_cmd(cmd, check=True, shell=True, capture_output=False):    try:        if capture_output:            result = subprocess.run(cmd, shell=shell, check=check,                                    stdout=subprocess.PIPE, stderr=subprocess.PIPE,                                    universal_newlines=True)            return result.stdout.strip()        else:            subprocess.run(cmd, shell=shell, check=check)    except subprocess.CalledProcessError as e:        print_colored(f"命令执行失败: {cmd}", Colors.RED)        if capture_output:            print_colored(f"错误信息: {e.stderr}", Colors.RED)        else:            print_colored(f"错误信息: {str(e)}", Colors.RED)        sys.exit(1)def get_ssh_version():    try:        version = run_cmd("ssh -V 2>&1", capture_output=True)        return version    except:        return "未安装或无法获取版本"def check_root():    if os.geteuid() != 0:        print_colored("错误：必须使用 root 用户执行此脚本！", Colors.RED)        sys.exit(1)def check_environment():    print_colored("\n[1/8] 正在检查系统环境...", Colors.BLUE)    redhat_release = Path("/etc/redhat-release")    if redhat_release.exists():        content = redhat_release.read_text()        if "CentOS Linux 7.9" not in content:            print_colored("警告：建议使用 CentOS 7.9，当前版本可能不完全兼容", Colors.YELLOW)    else:        print_colored("警告：未检测到 /etc/redhat-release，可能不是 RHEL/CentOS 系统", Colors.YELLOW)    arch = run_cmd("uname -m", capture_output=True)    if arch != "x86_64":        print_colored("错误：仅支持 64 位系统！", Colors.RED)        sys.exit(1)    print_colored("环境检查通过", Colors.GREEN)def smart_extract(tar_path, extract_to="/opt"):    tar_path = Path(tar_path)    if not tar_path.exists():        print_colored(f"错误：文件不存在 - {tar_path}", Colors.RED)        sys.exit(1)    if tar_path.stat().st_size == 0:        print_colored(f"错误：文件 {tar_path.name} 大小为0，可能未完整下载", Colors.RED)        sys.exit(1)    print_colored(f"解压 {tar_path.name} 到 {extract_to} ...", Colors.NC)    try:        with tarfile.open(tar_path, 'r:*') as tar:            tar.extractall(path=extract_to)            members = tar.getmembers()            if members:                top_dir = members[0].name.split('/')[0]                return top_dir            else:                return tar_path.name.replace('.tar.gz', '').replace('.tgz', '').replace('.tar', '')    except (tarfile.ReadError, EOFError) as e:        print_colored(f"Python tarfile 解压失败：{e}", Colors.RED)        print_colored("尝试使用系统 tar 命令解压...", Colors.YELLOW)        try:            subprocess.run(f"tar -xf {tar_path} -C {extract_to}", shell=True, check=True)            for item in Path(extract_to).iterdir():                if item.is_dir() and not item.name.startswith('.'):                    return item.name            return tar_path.name.replace('.tar.gz', '').replace('.tgz', '').replace('.tar', '')        except subprocess.CalledProcessError:            print_colored("系统 tar 命令也解压失败，文件可能严重损坏。", Colors.RED)            sys.exit(1)def install_dependencies(dep_file):    print_colored("\n[2/8] 安装基础依赖...", Colors.BLUE)    dep_path = Path(dep_file)    if not dep_path.exists():        print_colored(f"错误：依赖文件不存在: {dep_file}", Colors.RED)        skip = input("是否跳过依赖安装并继续？(y/n) [n]: ").strip().lower()        if skip in ('y', 'yes'):            print_colored("跳过依赖安装，请注意系统必须已具备编译环境！", Colors.YELLOW)            return        else:            sys.exit(1)    temp_dir = Path("/opt/deps_temp")    temp_dir.mkdir(exist_ok=True)    try:        top_dir_name = smart_extract(dep_file, extract_to=str(temp_dir))    except SystemExit:        skip = input("依赖包解压失败，是否跳过依赖安装并继续？(y/n) [n]: ").strip().lower()        if skip in ('y', 'yes'):            print_colored("跳过依赖安装，请注意系统必须已具备编译环境！", Colors.YELLOW)            shutil.rmtree(temp_dir, ignore_errors=True)            return        else:            shutil.rmtree(temp_dir, ignore_errors=True)            sys.exit(1)    extract_dir = temp_dir / top_dir_name    if not extract_dir.exists():        rpm_files = list(temp_dir.glob("**/*.rpm"))        if not rpm_files:            print_colored("错误：解压后未找到 RPM 包", Colors.RED)            shutil.rmtree(temp_dir)            sys.exit(1)        print_colored(f"找到 {len(rpm_files)} 个 RPM 包（分散存放），开始安装...", Colors.NC)        rpm_list = " ".join(f"'{f}'" for f in rpm_files)        run_cmd(f"rpm -ivh {rpm_list} --nodeps --force")    else:        os.chdir(extract_dir)        rpm_files = list(Path(".").glob("*.rpm"))        if not rpm_files:            print_colored("错误：yilai 目录下未找到任何 RPM 包", Colors.RED)            os.chdir("/opt")            shutil.rmtree(temp_dir)            sys.exit(1)        print_colored(f"在目录 '{top_dir_name}' 中找到 {len(rpm_files)} 个 RPM 包，开始安装...", Colors.NC)        run_cmd("rpm -ivh *.rpm --nodeps --force")        os.chdir("/opt")    shutil.rmtree(temp_dir)    print_colored("依赖包安装完成", Colors.GREEN)def build_zlib(zlib_file):    print_colored("\n[3/8] 编译安装 zlib...", Colors.BLUE)    zlib_path = Path(zlib_file)    if not zlib_path.exists():        print_colored(f"错误：zlib 源码包不存在: {zlib_file}", Colors.RED)        sys.exit(1)    dir_name = smart_extract(zlib_file)    extract_dir = Path("/opt") / dir_name    os.chdir(extract_dir)    run_cmd("./configure --prefix=/usr/local/zlib")    run_cmd("make")    run_cmd("make install")    with open("/etc/ld.so.conf", "a") as f:        f.write("/usr/local/zlib/lib\n")    run_cmd("ldconfig -v")    print_colored("zlib 安装完成", Colors.GREEN)def build_openssl(openssl_file):    print_colored("\n[4/8] 编译安装 OpenSSL...", Colors.BLUE)    ssl_path = Path(openssl_file)    if not ssl_path.exists():        print_colored(f"错误：OpenSSL 源码包不存在: {openssl_file}", Colors.RED)        sys.exit(1)    if "3." in ssl_path.name:        print_colored("警告：您正在使用 OpenSSL 3.x 版本，在 CentOS 7 下可能与 OpenSSH 存在链接问题。", Colors.YELLOW)        print_colored("强烈建议使用 OpenSSL 1.1.1 系列（如 openssl-1.1.1o.tar.gz）。", Colors.YELLOW)        cont = input("是否继续尝试？(y/n) [n]: ").strip().lower()        if cont not in ('y', 'yes'):            print_colored("用户取消操作，退出脚本。", Colors.YELLOW)            sys.exit(0)    dir_name = smart_extract(openssl_file)    extract_dir = Path("/opt") / dir_name    os.chdir(extract_dir)    run_cmd("make clean 2>/dev/null", check=False)    os.environ['CFLAGS'] = '-std=gnu99 -O2 -fPIC'    print_colored("已设置 CFLAGS=-std=gnu99 -O2 -fPIC", Colors.NC)    run_cmd("./config --prefix=/usr/local/ssl -d shared")    run_cmd("make")    run_cmd("make install")    with open("/etc/ld.so.conf", "a") as f:        f.write("/usr/local/ssl/lib\n")    run_cmd("ldconfig -v")    if not verify_openssl_install():        print_colored("OpenSSL 安装验证失败，请检查编译日志。", Colors.RED)        sys.exit(1)    print_colored("OpenSSL 安装完成并通过验证", Colors.GREEN)def verify_openssl_install():    libdir = Path("/usr/local/ssl/lib")    incdir = Path("/usr/local/ssl/include/openssl")    if not incdir.exists():        print_colored("错误：OpenSSL 头文件目录缺失", Colors.RED)        return False    libcrypto_so = libdir / "libcrypto.so"    libssl_so = libdir / "libssl.so"    libcrypto_a = libdir / "libcrypto.a"    libssl_a = libdir / "libssl.a"    if libcrypto_so.exists() and libssl_so.exists():        print_colored("OpenSSL 动态库 (.so) 存在", Colors.NC)        return True    elif libcrypto_a.exists() and libssl_a.exists():        print_colored("OpenSSL 静态库 (.a) 存在，将使用静态链接", Colors.YELLOW)        return True    else:        print_colored("错误：未找到 OpenSSL 库文件", Colors.RED)        return Falsedef get_openssl_version():    openssl_bin = Path("/usr/local/ssl/bin/openssl")    if openssl_bin.exists():        try:            version_str = subprocess.check_output([str(openssl_bin), "version"],                                                  universal_newlines=True).strip()            parts = version_str.split()            if len(parts) >= 2:                return parts[1]        except:            pass    return "unknown"def install_openssh(openssh_file):    print_colored("\n[5/8] 升级 OpenSSH...", Colors.BLUE)    ssh_path = Path(openssh_file)    if not ssh_path.exists():        print_colored(f"错误：OpenSSH 源码包不存在: {openssh_file}", Colors.RED)        sys.exit(1)    print_colored("卸载旧版本 OpenSSH...", Colors.NC)    run_cmd("rpm -e --nodeps openssh-server openssh openssh-clients 2>/dev/null", check=False)    dir_name = smart_extract(openssh_file)    extract_dir = Path("/opt") / dir_name    os.chdir(extract_dir)    sys_ssl_inc = Path("/usr/include/openssl")    sys_ssl_inc_backup = Path("/usr/include/openssl.bak")    if sys_ssl_inc.exists():        print_colored("备份系统 OpenSSL 头文件目录...", Colors.NC)        shutil.move(str(sys_ssl_inc), str(sys_ssl_inc_backup))    try:        run_cmd("ldconfig")        libdir = Path("/usr/local/ssl/lib")        use_static = False        if not (libdir / "libcrypto.so").exists():            use_static = True            print_colored("未找到动态库，将使用静态链接", Colors.YELLOW)        extra_cflags = "-I/usr/local/ssl/include"        ssl_version = get_openssl_version()        if ssl_version.startswith("3."):            extra_cflags += " -DOPENSSL_API_COMPAT=0x10100000L"            print_colored("已为 OpenSSL 3.x 添加兼容性宏", Colors.NC)        os.environ['CFLAGS'] = extra_cflags        os.environ['CPPFLAGS'] = extra_cflags        if use_static:            os.environ['LDFLAGS'] = f"-L{libdir}"            os.environ['LIBS'] = f"{libdir}/libcrypto.a {libdir}/libssl.a -ldl -lpthread"        else:            os.environ['LDFLAGS'] = f"-L{libdir} -Wl,-rpath,{libdir}"            os.environ['LIBS'] = "-lcrypto -lssl"            os.environ['LD_LIBRARY_PATH'] = str(libdir)        run_cmd("./configure --prefix=/usr/local/openssh "                "--with-zlib=/usr/local/zlib "                "--with-ssl-dir=/usr/local/ssl "                "--without-openssl-header-check")        run_cmd("make")        run_cmd("make install")    finally:        if sys_ssl_inc_backup.exists():            shutil.move(str(sys_ssl_inc_backup), str(sys_ssl_inc))    print_colored("配置 SSH 服务...", Colors.NC)    sshd_config = Path("/usr/local/openssh/etc/sshd_config")    with open(sshd_config, "a") as f:        f.write("\nPermitRootLogin yes\n")        f.write("PubkeyAuthentication yes\n")        f.write("PasswordAuthentication yes\n")    shutil.copy(sshd_config, "/etc/ssh/sshd_config")    with open("/etc/ssh/sshd_config", "a") as f:        f.write("HostKeyAlgorithms ssh-rsa,ssh-dss\n")    # 复制服务端二进制    sshd_bin = Path("/usr/sbin/sshd")    if sshd_bin.exists():        sshd_bin.rename("/usr/sbin/sshd.bak")    shutil.copy("/usr/local/openssh/sbin/sshd", "/usr/sbin/sshd")    os.chmod("/usr/sbin/sshd", 0o755)    # 复制客户端工具（包括 scp, sftp, ssh-add, ssh-agent, ssh-keyscan 等）    print_colored("安装客户端工具 (scp, sftp, ssh-keygen, ssh-add ...)", Colors.NC)    src_bin_dir = Path("/usr/local/openssh/bin")    dst_bin_dir = Path("/usr/bin")    tools_to_copy = [        "scp", "sftp", "ssh", "ssh-add", "ssh-agent", "ssh-keygen",        "ssh-keyscan", "sftp-server", "ssh-keysign"    ]    for tool in tools_to_copy:        src = src_bin_dir / tool        if src.exists():            dst = dst_bin_dir / tool            if dst.exists():                dst.rename(dst_bin_dir / f"{tool}.bak")            shutil.copy(src, dst)            os.chmod(dst, 0o755)            print_colored(f"  已安装: {tool}", Colors.NC)        else:            print_colored(f"  警告: 未找到 {tool}，跳过", Colors.YELLOW)    # 启动脚本    init_script = Path("/opt") / dir_name / "contrib/redhat/sshd.init"    if init_script.exists():        shutil.copy(init_script, "/etc/init.d/sshd")        os.chmod("/etc/init.d/sshd", 0o755)        run_cmd("chkconfig --add sshd")        run_cmd("chkconfig sshd on")    print_colored("OpenSSH 安装完成，所有客户端工具已就绪", Colors.GREEN)def final_check(old_version=""):    print_colored("\n[6/8] 执行最终检查...", Colors.BLUE)    run_cmd("systemctl daemon-reload")    run_cmd("systemctl restart sshd")    new_version = get_ssh_version()    if "OpenSSH_" in new_version:        print_colored("=" * 50, Colors.GREEN)        print_colored("升级成功！", Colors.GREEN)        if old_version:            print_colored(f"升级前版本：{old_version}", Colors.YELLOW)        print_colored(f"升级后版本：{new_version}", Colors.GREEN)        print_colored("=" * 50, Colors.GREEN)        print_colored("警告：请通过新 SSH 端口连接确认无误后，再关闭 Telnet 服务！", Colors.YELLOW)        # 验证 scp 是否可用        if Path("/usr/bin/scp").exists():            print_colored("scp 命令已安装，可以使用。", Colors.GREEN)    else:        print_colored(f"错误：升级失败，当前版本: {new_version}", Colors.RED)        sys.exit(1)def prompt_file(prompt_text, default_path=None, allow_skip=False):    while True:        if default_path:            user_input = input(f"{prompt_text} [默认: {default_path}]: ").strip()            if not user_input:                user_input = default_path        else:            user_input = input(f"{prompt_text}: ").strip()        if allow_skip and user_input.lower() in ('skip', 's', ''):            return None        if not user_input:            print_colored("路径不能为空，请重新输入。", Colors.YELLOW)            continue        if Path(user_input).exists():            return user_input        else:            print_colored(f"文件不存在: {user_input}", Colors.RED)            choice = input("是否重新输入？(y/n) [y]: ").strip().lower()            if choice in ('n', 'no'):                print_colored("用户取消操作，退出脚本。", Colors.YELLOW)                sys.exit(0)def main():    parser = argparse.ArgumentParser(description="OpenSSH 离线升级工具")    parser.add_argument("--dep", help="依赖包完整路径")    parser.add_argument("--zlib", help="zlib 源码包完整路径")    parser.add_argument("--openssl", help="OpenSSL 源码包完整路径")    parser.add_argument("--openssh", help="OpenSSH 源码包完整路径")    parser.add_argument("--skip-env-check", action="store_true", help="跳过环境检查")    args = parser.parse_args()    print_colored("========================================", Colors.BLUE)    print_colored("      OpenSSH 离线升级工具 (Python版)    ", Colors.GREEN)    print_colored("========================================", Colors.BLUE)    check_root()    old_version = get_ssh_version()    print_colored(f"\n当前系统 SSH 版本：{old_version}", Colors.YELLOW)    if not args.skip_env_check:        check_environment()    else:        print_colored("警告：已跳过环境检查", Colors.YELLOW)    print_colored("\n>>> 请准备以下四个文件（必须全部存在）：", Colors.YELLOW)    print_colored("    1. 依赖包 (yilai.tar.gz)  - 包含编译所需的 RPM 包", Colors.NC)    print_colored("    2. zlib 源码包 (如 zlib-1.3.1.tar.gz)", Colors.NC)    print_colored("    3. OpenSSL 源码包 (强烈建议使用 openssl-1.1.1o.tar.gz)", Colors.NC)    print_colored("    4. OpenSSH 源码包 (如 openssh-9.9p1.tar.gz)", Colors.NC)    dep_file = args.dep or prompt_file("请输入依赖包路径", default_path="/opt/yilai.tar.gz")    zlib_file = args.zlib or prompt_file("请输入 zlib 源码包路径", default_path="/opt/zlib-1.3.1.tar.gz")    openssl_file = args.openssl or prompt_file("请输入 OpenSSL 源码包路径", default_path="/opt/openssl-1.1.1o.tar.gz")    openssh_file = args.openssh or prompt_file("请输入 OpenSSH 源码包路径", default_path="/opt/openssh-9.9p1.tar.gz")    print_colored("\n========================================", Colors.BLUE)    print_colored("即将使用以下文件进行升级：", Colors.GREEN)    print_colored(f"依赖包路径: {dep_file}", Colors.NC)    print_colored(f"zlib 包路径: {zlib_file}", Colors.NC)    print_colored(f"OpenSSL 包路径: {openssl_file}", Colors.NC)    print_colored(f"OpenSSH 包路径: {openssh_file}", Colors.NC)    print_colored("========================================", Colors.BLUE)    confirm = input("确认开始升级？(y/n) [y]: ").strip().lower()    if confirm not in ('', 'y', 'yes'):        print_colored("用户取消操作，退出脚本。", Colors.YELLOW)        sys.exit(0)    install_dependencies(dep_file)    build_zlib(zlib_file)    build_openssl(openssl_file)    install_openssh(openssh_file)    final_check(old_version)if __name__ == "__main__":    main()