wecom_it_smart_desk/docs/09-部署运维/toolbox/fast_upload.py

#!/usr/bin/env python3
"""
快速 base64 上传脚本 — 大分块版本
将本地文件通过 JumpServer PTY base64 通道上传到远程服务器
使用 8000 字符/块（远大于 jms_ops.py 的 500 字符），大幅减少命令数
"""
import sys
import base64
import hashlib
from pathlib import Path

# 添加 jms_ops.py 所在目录
SKILL_DIR = Path(r"C:\Users\simon\.workbuddy\skills\jumpserver-ops\scripts")
sys.path.insert(0, str(SKILL_DIR))

# 导入 jms_ops 中的核心函数
from jms_ops import get_connection_tokens, PlinkSession


def fast_upload(local_path: str, remote_path: str, chunk_size: int = 8000):
    """大分块 base64 上传"""
    local_file = Path(local_path)
    if not local_file.exists():
        print(f"ERROR: file not found: {local_path}")
        return False

    local_data = local_file.read_bytes()
    local_md5 = hashlib.md5(local_data).hexdigest()
    b64_data = base64.b64encode(local_data).decode("ascii")

    # 分块
    chunks = [b64_data[i:i+chunk_size] for i in range(0, len(b64_data), chunk_size)]
    total_chunks = len(chunks)

    print(f"File: {local_file.name}")
    print(f"Size: {len(local_data)} bytes")
    print(f"Base64: {len(b64_data)} chars")
    print(f"Chunks: {total_chunks} x {chunk_size} chars")
    print(f"MD5: {local_md5}")
    print(f"Target: {remote_path}")
    print()

    # 获取 token + 启动会话
    tokens = get_connection_tokens(1)
    if not tokens:
        print("ERROR: failed to get connection tokens")
        return False

    token_id, token_secret = tokens[0]
    session = PlinkSession(f"JMS-{token_id}", token_secret)
    if not session.connect():
        print("ERROR: failed to connect session")
        return False

    try:
        # 1. 清空目标文件
        print("Clearing target file...")
        session.run_command(f"> {remote_path}", timeout=5)

        # 2. 逐块追加
        for i, chunk in enumerate(chunks):
            # 用 printf 避免 echo 的换行符问题
            cmd = f"printf '%s' '{chunk}' >> {remote_path}.b64"
            r = session.run_command(cmd, timeout=10)
            if not r["success"]:
                print(f"  FAIL chunk {i+1}/{total_chunks}")
                return False

            # 进度报告
            if (i+1) % 20 == 0 or (i+1) == total_chunks:
                pct = (i+1) * 100 // total_chunks
                print(f"  [{pct:3d}%] chunk {i+1}/{total_chunks}")

        # 3. base64 解码
        print(f"\nDecoding base64 -> {remote_path}...")
        r = session.run_command(f"base64 -d {remote_path}.b64 > {remote_path}", timeout=30)
        if not r["success"]:
            print(f"  WARN decode result: {r}")

        # 4. 验证大小
        print("Verifying size...")
        r = session.run_command(f"wc -c < {remote_path}", timeout=5)
        if r["success"]:
            remote_size_str = r["output"].strip()
            remote_size = int(remote_size_str) if remote_size_str.isdigit() else -1
            if remote_size == len(local_data):
                print(f"  OK size match: {remote_size} bytes")
            else:
                print(f"  SIZE MISMATCH: local={len(local_data)}, remote={remote_size}")
                return False
        else:
            print(f"  WARN cannot verify size: {r}")
            return True  # 仍然认为成功

        # 5. MD5 验证
        print("Verifying MD5...")
        r = session.run_command(f"md5sum {remote_path}", timeout=10)
        if r["success"]:
            remote_md5 = r["output"].split()[0]
            if remote_md5 == local_md5:
                print(f"  OK MD5 match: {remote_md5}")
            else:
                print(f"  MD5 MISMATCH: local={local_md5}, remote={remote_md5}")
                # 大小匹配但 MD5 不匹配，可能是 PTY 换行符问题
                print("  (size matches, trying gzip test instead)")
                r2 = session.run_command(f"gzip -t {remote_path} 2>&1 && echo GZIP_OK || echo GZIP_FAIL", timeout=10)
                if r2["success"] and "GZIP_OK" in r2["output"]:
                    print("  OK gzip integrity test passed")
                    # 清理临时文件
                    session.run_command(f"rm -f {remote_path}.b64", timeout=5)
                    return True
                else:
                    print(f"  GZIP FAIL: {r2}")
                    return False
        else:
            print(f"  WARN cannot verify MD5")

        # 6. 清理临时文件
        session.run_command(f"rm -f {remote_path}.b64", timeout=5)
        print("\nUpload complete!")
        return True

    finally:
        session.close()


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Fast base64 upload via JumpServer PTY")
    parser.add_argument("local", help="Local file path")
    parser.add_argument("remote", help="Remote file path")
    parser.add_argument("--chunk-size", type=int, default=8000, help="Chunk size in chars (default: 8000)")
    args = parser.parse_args()

    success = fast_upload(args.local, args.remote, args.chunk_size)
    sys.exit(0 if success else 1)