【跨服务器的数据自动化下载--安装公钥，免密下载】

发现好久没写csdn了，说多了都是泪~~
以后会更新一些自动化工作的脚本or 小tricks，欢迎交流。分享一个最近在业务上写的较为实用的自动化脚本，可以批量从远端服务器下载指定数据到当前程序运行服务器，无需每次都输入密码，通过集成安装公钥的功能，实现免密下载。这个方式适合这种不能安装 sshpass又想全自动化的环境。

功能介绍：

第一次运行：会提示输入一次远端 服务器 的密码（安装公钥）。安装完成后立即用免密下载。
以后运行：直接免密下载，全自动。
后续所有下载脚本全自动化，不用再输入任何密码，也不依赖额外工具。
以从远端服务器（账户：abc）下载数据为例。

上代码：

#!/usr/bin/env python3
import sys
import os
import subprocess
import tempfile
import shutil

def check_key_auth(remote_username, remote_host, private_key_path):
    """检查是否已配置免密"""
    result = subprocess.run(
        ['ssh', '-i', private_key_path, '-o', 'BatchMode=yes', '-o', 'StrictHostKeyChecking=no',
         f'{remote_username}@{remote_host}', 'echo ok'],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    return result.returncode == 0

def install_pubkey(remote_username, remote_host, private_key_path):
    """将本地公钥写入远端 authorized_keys（会要求输入密码一次）"""
    pubkey_path = private_key_path + '.pub'
    if not os.path.exists(pubkey_path):
        print(f"未找到公钥文件 {pubkey_path}")
        return False

    with open(pubkey_path, 'r') as f:
        pubkey_content = f.read().strip()

    print(f"正在将公钥写入 {remote_host}，需要输入 {remote_username} 的密码...")
    cmd = [
        'ssh', f'{remote_username}@{remote_host}',
        f'mkdir -p ~/.ssh && chmod 700 ~/.ssh && echo "{pubkey_content}" >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys'
    ]
    return subprocess.run(cmd).returncode == 0

def download_files_batch(remote_username, remote_host, private_key_path, files_to_download):
    """
    批量下载实现：
      1. 在本地创建一个临时目录 tmpdir
      2. 用 single scp 命令把所有 remote:files 下载到 tmpdir（scp 支持多个 source + 单个 destination directory）
      3. 将 tmpdir 下的文件移动到各自的目标位置
    """
    # 先保证目标本地目录存在（创建父目录）
    for _, local_path in files_to_download:
        os.makedirs(os.path.dirname(local_path), exist_ok=True)

    tmpdir = tempfile.mkdtemp(prefix='download_')
    try:
        # 构造 scp 命令：scp -i key user@host:/path/to/file1 user@host:/path/to/file2 ... <tmpdir>
        scp_cmd = ['scp', '-i', private_key_path, '-o', 'StrictHostKeyChecking=no']
        for remote_path, _ in files_to_download:
            scp_cmd.append(f'{remote_username}@{remote_host}:{remote_path}')
        scp_cmd.append(tmpdir)  # scp 要求最后是目标目录

        print("执行 scp，目标临时目录：", tmpdir)
        print("scp 命令：", ' '.join(scp_cmd))
        proc = subprocess.run(scp_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if proc.returncode != 0:
            print("批量 scp 失败，错误信息：")
            print(proc.stderr.decode(errors='ignore'))
            return False

        # scp 成功，把文件从 tmpdir 移动到最终位置
        for remote_path, local_path in files_to_download:
            fname = os.path.basename(remote_path)
            src = os.path.join(tmpdir, fname)
            if not os.path.exists(src):
                print(f"警告：远端文件 {remote_path} 未被下载到临时目录（{src} 不存在）")
                # 继续处理其它文件
                continue
            shutil.move(src, local_path)
            print(f"已移动：{src} -> {local_path}")

        print("全部文件处理完成。")
        return True

    finally:
        # 清理临时目录（如果里面还有剩余文件，会一并删除）
        try:
            shutil.rmtree(tmpdir)
        except Exception as e:
            print("清理临时目录时出错：", e)

def main():
    if len(sys.argv) != 2:
        print("用法: python fully_auto_download.py <date>")
        sys.exit(1)

    date = sys.argv[1]
    remote_host = "" ##你需要获取数据的远端ip
    remote_username = "abc"
    private_key_path = '/.ssh/id_rsa'

    files_to_download = [
        (
            f"/data/{date[:4]}/{date}_000000_000000.csv",
            f"/save/{date}_000000_000000.csv"
        ),
        ## 此处省略其他需要下载的数据
    ]

    # 检查免密
    if not check_key_auth(remote_username, remote_host, private_key_path):
        print("未检测到免密，开始安装公钥...")
        if not install_pubkey(remote_username, remote_host, private_key_path):
            print("公钥安装失败，请检查密码是否正确或远端权限设置。")
            sys.exit(1)
        if not check_key_auth(remote_username, remote_host, private_key_path):
            print("免密配置仍然失败，请手动检查。")
            sys.exit(1)
        print("免密配置成功！")

    # 一次性批量下载
    ok = download_files_batch(remote_username, remote_host, private_key_path, files_to_download)
    if not ok:
        print("批量下载出现问题，退回逐个下载尝试。")
        # 如果批量下载失败，可以回退到逐个 scp（免密已经配置好了，不会再要求密码）
        for remote_path, local_path in files_to_download:
            try:
                os.makedirs(os.path.dirname(local_path), exist_ok=True)
                scp_cmd = [
                    'scp', '-i', private_key_path, '-o', 'StrictHostKeyChecking=no',
                    f'{remote_username}@{remote_host}:{remote_path}', local_path
                ]
                subprocess.run(scp_cmd, check=True)
                print(f"下载成功: {local_path}")
            except subprocess.CalledProcessError as e:
                print(f"下载失败: {local_path}, 错误: {e}")

if __name__ == "__main__":
    main()