soul-yongping/scripts/pull_from_baota.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从宝塔正式机拉取线上运行目录到本地镜像（与 soul-api/master.py、soul-admin/master.py 同源 SSH 配置）。

说明：
- 服务器上一般是「二进制 + .env + 日志」与「静态 dist」，不包含完整 Go/React 源码。
- 默认解压到仓库根目录 _server_live/soul-api、_server_live/soul-admin，不覆盖本地工程源码。

环境变量与 master.py 一致：DEPLOY_HOST、DEPLOY_USER、DEPLOY_PASSWORD、DEPLOY_SSH_KEY、
DEPLOY_PROJECT_PATH、DEPLOY_BASE_PATH。
"""

from __future__ import print_function

import argparse
import importlib.util
import os
import shutil
import sys
import tarfile
import tempfile
import threading

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
SOUL_API_DIR = os.path.join(ROOT, "soul-api")


def _load_api_master():
    path = os.path.join(SOUL_API_DIR, "master.py")
    spec = importlib.util.spec_from_file_location("soul_api_deploy_master", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def _pull_dir_tar(client, remote_dir, local_dir, mod, timeout=600):
    """远端 tar czf 流式下载并解压到 local_dir。"""
    import shlex

    remote_q = shlex.quote(remote_dir)
    cmd = "tar czf - -C %s . 2>/dev/null" % remote_q
    stdin, stdout, stderr = client.exec_command(cmd, timeout=timeout)

    err_holder = []

    def _drain():
        try:
            err_holder.append(stderr.read())
        except Exception:
            pass

    t = threading.Thread(target=_drain)
    t.daemon = True
    t.start()

    fd, tmp_path = tempfile.mkstemp(suffix=".tar.gz")
    os.close(fd)
    try:
        with open(tmp_path, "wb") as out:
            while True:
                chunk = stdout.read(256 * 1024)
                if not chunk:
                    break
                out.write(chunk)
        t.join(timeout=5)
        exit_status = stdout.channel.recv_exit_status()
        if exit_status != 0:
            print("  [警告] 远端 tar 退出码: %s" % exit_status)
        if os.path.isdir(local_dir):
            shutil.rmtree(local_dir)
        os.makedirs(local_dir, exist_ok=True)
        with tarfile.open(tmp_path, "r:gz") as tf:
            tf.extractall(local_dir)
        print("  [成功] 已同步到: %s" % local_dir)
        return True
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass


def main():
    parser = argparse.ArgumentParser(description="从宝塔拉取 soul-api / soul-admin 线上目录")
    parser.add_argument("--api-only", action="store_true", help="仅拉 soul-api")
    parser.add_argument("--admin-only", action="store_true", help="仅拉 soul-admin")
    args = parser.parse_args()

    mod = _load_api_master()
    cfg = mod.get_cfg()
    if not cfg.get("password") and not cfg.get("ssh_key"):
        print("[失败] 请设置 DEPLOY_PASSWORD 或 DEPLOY_SSH_KEY")
        return 1

    pull_api = not args.admin_only
    pull_admin = not args.api_only
    if args.api_only and args.admin_only:
        print("[失败] 不能同时指定 --api-only 与 --admin-only")
        return 1

    client = None
    try:
        client = mod._connect_ssh(cfg)
        live_root = os.path.join(ROOT, "_server_live")
        os.makedirs(live_root, exist_ok=True)

        print("=" * 60)
        print("  从宝塔拉取线上目录 → %s" % live_root)
        print("  主机: %s@%s:%s" % (cfg["user"], cfg["host"], mod.DEFAULT_SSH_PORT))
        print("=" * 60)

        if pull_api:
            print("[1] soul-api: %s" % cfg["project_path"])
            _pull_dir_tar(
                client,
                cfg["project_path"],
                os.path.join(live_root, "soul-api"),
                mod,
            )

        if pull_admin:
            admin_base = os.environ.get("DEPLOY_BASE_PATH", "/www/wwwroot/self/soul-admin").rstrip("/")
            print("[2] soul-admin: %s" % admin_base)
            # 复用同一连接；若仅拉 admin，上面未开新连接也行
            if not pull_api:
                pass
            _pull_dir_tar(
                client,
                admin_base,
                os.path.join(live_root, "soul-admin"),
                mod,
            )

        print("")
        print("  完成。镜像根目录: %s" % live_root)
        return 0
    except Exception as e:
        print("[失败] %s" % e)
        import traceback

        traceback.print_exc()
        return 1
    finally:
        if client:
            try:
                client.close()
            except Exception:
                pass


if __name__ == "__main__":
    sys.exit(main() or 0)