#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""远程压测机器操作脚本（远程机器配置 + SSH 执行 + 文件推拉 + 远程 aicore 扫描）

配合 config/remote_machines.json 使用（模板见 config/remote_machines.example.json）。
支持配置多台远程机器（ip / ssh_port / serial_port 串口号 / username / password 等），
批量压测时可将一部分接口分配到远程机器上压测（详见 AGENTS.md "远程机器配置与调度" 章节）。

子命令:
  list    列出已配置的远程机器（密码脱敏）
  check   远程机器连通性与环境检查（SSH / npu-smi 卡数 / CANN 环境 / 磁盘 / remote_root）
  exec    在指定远程机器上执行命令（长时压测命令需配 --timeout 秒）
  push    推送本地文件/目录到远程机器（turing_sec / output CPP / scripts 环境准备）
  pull    拉取远程文件/目录到本地（打屏日志 / plog 回传 output/{aclnnApiName}/fuzz_logs/）
  scan    远程 aicore_error 扫描（远程 plog 目录 + 远程 output 目录 grep -rn "kernel_name"）

SSH 实现: 优先 paramiko（跨平台、密码免交互）；paramiko 未安装时回退 sshpass+ssh/scp/rsync（Linux）。

安全规范:
  - config/remote_machines.json 已被 .gitignore 排除，严禁提交 git
  - password 支持 ${ENV_VAR} 环境变量引用（推荐，避免明文）
  - 任何输出（日志/报错）严禁回显密码

用法示例:
  python scripts/remote_ops.py list
  python scripts/remote_ops.py check                      # 检查全部 enabled 机器
  python scripts/remote_ops.py check --name remote-1      # 检查指定机器
  python scripts/remote_ops.py exec  --name remote-1 --cmd "npu-smi info"
  python scripts/remote_ops.py exec  --name remote-1 --cmd "cd ~/fuzz_remote/turing_sec/aclnnAPI_fuzz/out && ./aclnnAPI_fuzz aclnnAbs_fuzz 10000" --timeout 1800
  python scripts/remote_ops.py push  --name remote-1 --local output/aclnnAbs --remote ~/fuzz_remote/output/aclnnAbs
  python scripts/remote_ops.py pull  --name remote-1 --remote "~/fuzz_remote/output/aclnnAbs/fuzz_logs" --local output/aclnnAbs
  python scripts/remote_ops.py scan  --name remote-1

退出码: 0=成功; 1=运行错误; 2=check/scan 发现异常（连接失败/aicore_error 命中）
"""
import argparse
import json
import os
import re
import shutil
import stat as stat_mod
import subprocess
import sys
import time

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

DEFAULT_CONFIG = os.path.join("config", "remote_machines.json")
DEFAULT_PLOG_DIR = "/root/ascend/log/debug/plog"

AICORE_ERROR_PATTERNS = [
    "the error is aivec error",
    "there is an aivec error exception",
    "there is an aicore error exception",
]


def log(tag, msg):
    print("[{}] {}".format(tag, msg))


def try_import_paramiko():
    try:
        import paramiko

        return paramiko
    except ImportError:
        return None


class MachineError(Exception):
    pass


def resolve_secret(value, field, machine_name):
    """解析密码字段：支持 ${ENV_VAR} 环境变量引用，避免明文存储。"""
    if value is None:
        raise MachineError("机器 {} 缺少必填字段 {}".format(machine_name, field))
    if value.startswith("${") and value.endswith("}"):
        env_name = value[2:-1]
        real = os.environ.get(env_name)
        if not real:
            raise MachineError(
                "机器 {} 的 {} 引用环境变量 {} 未设置（export {}=... 后重试）".format(
                    machine_name, field, env_name, env_name))
        return real
    return value


class RemoteMachine:
    """单台远程机器配置（ip / ssh_port / serial_port 串口号 / username / password 等）。

    password 惰性解析：仅在首次连接时解析 ${ENV_VAR} 引用，避免禁用机器/未设置环境变量
    影响其他机器的加载与使用。
    """

    def __init__(self, raw, config_path):
        self.config_path = config_path
        self.name = str(raw.get("name") or "").strip()
        if not self.name:
            raise MachineError("配置 {} 中存在缺少 name 的机器条目".format(config_path))
        self.ip = str(raw.get("ip") or "").strip()
        if not self.ip:
            raise MachineError("机器 {} 缺少必填字段 ip".format(self.name))
        self.ssh_port = int(raw.get("ssh_port") or 22)
        self.serial_port = str(raw.get("serial_port") or "").strip()  # 串口号（BMC 串口管理/标识用）
        self.username = str(raw.get("username") or "root").strip()
        self._password_raw = raw.get("password")
        self._password_resolved = None
        self.remote_root = raw.get("remote_root") or "~/fuzz_remote"
        self.max_parallel = max(1, int(raw.get("max_parallel") or 1))
        self.enabled = bool(raw.get("enabled", True))
        self.notes = raw.get("notes") or ""

    @property
    def password(self):
        if self._password_resolved is None:
            self._password_resolved = resolve_secret(
                self._password_raw, "password", self.name)
        return self._password_resolved

    def summary(self):
        return "name={} enabled={} ip={} ssh_port={} serial_port={} user={} remote_root={} max_parallel={}".format(
            self.name, self.enabled, self.ip, self.ssh_port,
            self.serial_port or "-", self.username, self.remote_root, self.max_parallel)

    # ---------------- paramiko 通道 ----------------
    def _connect(self, paramiko, timeout=15):
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(
            self.ip, port=self.ssh_port, username=self.username, password=self.password,
            timeout=timeout, allow_agent=False, look_for_keys=False,
        )
        return client

    # ---------------- 命令执行 ----------------
    def exec_cmd(self, cmd, timeout=600, silent=False):
        """远程执行命令，返回 (exit_code, combined_output)。优先 paramiko，回退 sshpass。"""
        paramiko = try_import_paramiko()
        if paramiko is not None:
            return self._exec_paramiko(paramiko, cmd, timeout, silent)
        return self._exec_cli(cmd, timeout, silent)

    def _exec_paramiko(self, paramiko, cmd, timeout, silent):
        client = None
        try:
            client = self._connect(paramiko)
            transport = client.get_transport()
            channel = transport.open_session()
            channel.set_combine_stderr(True)
            channel.exec_command(cmd)
            deadline = time.time() + max(1, timeout)
            chunks = []
            while True:
                got = False
                while channel.recv_ready():
                    data = channel.recv(65536).decode("utf-8", errors="replace")
                    chunks.append(data)
                    if not silent:
                        sys.stdout.write(data)
                    got = True
                if channel.exit_status_ready() and not channel.recv_ready():
                    break
                if time.time() > deadline:
                    channel.close()
                    raise MachineError(
                        "机器 {} 命令执行超时({}s)，命令已被终止: {}".format(self.name, timeout, cmd[:120]))
                if not got:
                    time.sleep(0.05)
            code = channel.recv_exit_status()
            return code, "".join(chunks)
        finally:
            if client is not None:
                try:
                    client.close()
                except Exception:
                    pass

    def _exec_cli(self, cmd, timeout, silent):
        if shutil.which("sshpass") is None:
            raise MachineError(
                "机器 {} 无法执行远程命令：paramiko 未安装且系统无 sshpass。"
                "请执行 pip install paramiko（推荐，跨平台）或在 Linux 上安装 sshpass".format(self.name))
        argv = [
            "sshpass", "-p", self.password,
            "ssh", "-p", str(self.ssh_port),
            "-o", "StrictHostKeyChecking=no",
            "-o", "UserKnownHostsFile=/dev/null",
            "-o", "ConnectTimeout=15",
            "{}@{}".format(self.username, self.ip), cmd,
        ]
        try:
            proc = subprocess.run(argv, capture_output=True, text=True,
                                  encoding="utf-8", errors="replace", timeout=timeout)
        except subprocess.TimeoutExpired:
            raise MachineError("机器 {} 命令执行超时({}s): {}".format(self.name, timeout, cmd[:120]))
        output = (proc.stdout or "") + (proc.stderr or "")
        if not silent:
            sys.stdout.write(output)
        return proc.returncode, output

    # ---------------- 文件传输 ----------------
    def push(self, local, remote, timeout=1800):
        """推送本地文件/目录到远程（目录递归）。优先 rsync（增量高效），回退 paramiko SFTP / scp。"""
        if not os.path.exists(local):
            raise MachineError("本地路径不存在: {}".format(local))
        if shutil.which("rsync") is not None:
            argv = [
                "rsync", "-az", "--stats",
                "-e", "ssh -p {} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null".format(self.ssh_port),
                local.rstrip("/\\") + "/", remote.rstrip("/") + "/",
            ]
            env = dict(os.environ)
            if shutil.which("sshpass") is not None:
                argv = ["sshpass", "-p", self.password] + argv
            else:
                # rsync 无 sshpass 时依赖 SSH 免密（密钥）；失败则回退 paramiko
                proc = subprocess.run(argv, capture_output=True, text=True,
                                      encoding="utf-8", errors="replace", timeout=timeout, env=env)
                if proc.returncode == 0:
                    log("PUSH", "{} -> {}@{}:{} (rsync) OK".format(
                        local, self.username, self.ip, remote))
                    return
        paramiko = try_import_paramiko()
        if paramiko is not None:
            self._push_sftp(paramiko, local, remote)
            return
        if shutil.which("sshpass") is not None and shutil.which("scp") is not None:
            argv = ["sshpass", "-p", self.password, "scp", "-r", "-P", str(self.ssh_port),
                    "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
                    local, "{}@{}:{}".format(self.username, self.ip, remote)]
            proc = subprocess.run(argv, capture_output=True, text=True,
                                  encoding="utf-8", errors="replace", timeout=timeout)
            if proc.returncode != 0:
                raise MachineError("机器 {} scp 推送失败: {}".format(self.name, (proc.stderr or "")[:200]))
            log("PUSH", "{} -> {}@{}:{} (scp) OK".format(local, self.username, self.ip, remote))
            return
        raise MachineError("机器 {} 无可用传输通道：请 pip install paramiko 或安装 rsync/scp+sshpass".format(self.name))

    def _push_sftp(self, paramiko, local, remote):
        client = None
        try:
            client = self._connect(paramiko)
            sftp = client.open_sftp()
            self._sftp_mkdirs(sftp, os.path.dirname(remote.rstrip("/")) or ".")
            self._push_walk(sftp, local, remote)
            log("PUSH", "{} -> {}@{}:{} (sftp) OK".format(local, self.username, self.ip, remote))
        finally:
            if client is not None:
                try:
                    client.close()
                except Exception:
                    pass

    def _push_walk(self, sftp, local, remote):
        if os.path.isfile(local):
            sftp.put(local, remote)
            return
        try:
            sftp.mkdir(remote)
        except IOError:
            pass  # 已存在
        for item in sorted(os.listdir(local)):
            if item in (".git", "__pycache__", "out", "dumpdata"):
                # 跳过 git 元数据/缓存/编译产物，减小同步体积（out 可按需单独推）
                continue
            self._push_walk(sftp, os.path.join(local, item), remote.rstrip("/") + "/" + item)

    def _sftp_mkdirs(self, sftp, remote_dir):
        remote_dir = remote_dir.rstrip("/")
        if remote_dir in ("", ".", "..", "/"):
            return
        try:
            sftp.stat(remote_dir)
            return
        except IOError:
            pass
        parent = remote_dir.rsplit("/", 1)[0] if "/" in remote_dir else ""
        if parent and parent != remote_dir:
            self._sftp_mkdirs(sftp, parent)
        try:
            sftp.mkdir(remote_dir)
        except IOError:
            pass

    def pull(self, remote, local, timeout=1800):
        """拉取远程文件/目录到本地（目录递归）。优先 rsync，回退 paramiko SFTP / scp。"""
        if shutil.which("rsync") is not None:
            argv = [
                "rsync", "-az",
                "-e", "ssh -p {} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null".format(self.ssh_port),
                "{}@{}:{}".format(self.username, self.ip, remote.rstrip("/")) + "/", local.rstrip("/\\") + "/",
            ]
            if shutil.which("sshpass") is not None:
                argv = ["sshpass", "-p", self.password] + argv
                try:
                    proc = subprocess.run(argv, capture_output=True, text=True,
                                          encoding="utf-8", errors="replace", timeout=timeout)
                    if proc.returncode == 0:
                        log("PULL", "{}@{}:{} -> {} (rsync) OK".format(self.username, self.ip, remote, local))
                        return
                except subprocess.TimeoutExpired:
                    raise MachineError("机器 {} rsync 拉取超时({}s)".format(self.name, timeout))
        paramiko = try_import_paramiko()
        if paramiko is not None:
            self._pull_sftp(paramiko, remote, local)
            return
        if shutil.which("sshpass") is not None and shutil.which("scp") is not None:
            os.makedirs(local, exist_ok=True)
            argv = ["sshpass", "-p", self.password, "scp", "-r", "-P", str(self.ssh_port),
                    "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
                    "{}@{}:{}".format(self.username, self.ip, remote), local]
            proc = subprocess.run(argv, capture_output=True, text=True,
                                  encoding="utf-8", errors="replace", timeout=timeout)
            if proc.returncode != 0:
                raise MachineError("机器 {} scp 拉取失败: {}".format(self.name, (proc.stderr or "")[:200]))
            log("PULL", "{}@{}:{} -> {} (scp) OK".format(self.username, self.ip, remote, local))
            return
        raise MachineError("机器 {} 无可用传输通道：请 pip install paramiko 或安装 rsync/scp+sshpass".format(self.name))

    def _pull_sftp(self, paramiko, remote, local):
        client = None
        try:
            client = self._connect(paramiko)
            sftp = client.open_sftp()
            self._pull_walk(sftp, remote, local)
            log("PULL", "{}@{}:{} -> {} (sftp) OK".format(self.username, self.ip, remote, local))
        finally:
            if client is not None:
                try:
                    client.close()
                except Exception:
                    pass

    def _pull_walk(self, sftp, remote, local):
        st = sftp.lstat(remote)
        if stat_mod.S_ISREG(st.st_mode):
            os.makedirs(os.path.dirname(os.path.abspath(local)), exist_ok=True)
            sftp.get(remote, local)
            return
        os.makedirs(local, exist_ok=True)
        for item in sorted(sftp.listdir(remote)):
            self._pull_walk(sftp, remote.rstrip("/") + "/" + item,
                            os.path.join(local, item))


def load_machines(config_path, name_filter=None, only_enabled=True):
    if not os.path.isfile(config_path):
        raise MachineError(
            "远程机器配置不存在: {}\n"
            "请复制模板 config/remote_machines.example.json 为 config/remote_machines.json 并填写"
            "（ip/ssh_port/serial_port 串口号/username/password，支持多台机器）".format(config_path))
    with open(config_path, "r", encoding="utf-8") as fh:
        data = json.load(fh)
    machines_raw = data.get("machines") if isinstance(data, dict) else data
    if not isinstance(machines_raw, list):
        raise MachineError("配置 {} 格式错误：应为 {{\"machines\": [...]}} 或机器数组".format(config_path))
    # 过滤在构造前完成：name 查找不受其他机器配置问题影响；only_enabled 跳过禁用机器
    if name_filter:
        machines_raw = [r for r in machines_raw
                        if str(r.get("name") or "").strip() == name_filter]
        if not machines_raw:
            raise MachineError("配置 {} 中未找到名为 {} 的机器".format(config_path, name_filter))
    elif only_enabled:
        machines_raw = [r for r in machines_raw if r.get("enabled", True)]
    return [RemoteMachine(raw, config_path) for raw in machines_raw]


def expand_remote_path(machine, path):
    """展开远程路径中的 ~ 与 $REMOTE_ROOT。"""
    path = path.replace("$REMOTE_ROOT", machine.remote_root)
    if path.startswith("~"):
        # 远端 shell 展开 ~
        code, out = machine.exec_cmd("echo {}".format(path), timeout=30, silent=True)
        if code == 0:
            return out.strip().splitlines()[-1] if out.strip() else path
    return path


# ---------------- 子命令 ----------------

def cmd_list(args):
    machines = load_machines(args.config, only_enabled=False)
    log("LIST", "共配置 {} 台远程机器（{}），配置文件: {}".format(
        len(machines), "密码已脱敏", args.config))
    for m in machines:
        print("  - {}".format(m.summary()))
        if m.notes:
            print("      notes: {}".format(m.notes))
    return 0


def cmd_check(args):
    machines = load_machines(args.config, args.name)
    if not machines:
        log("CHECK", "无 enabled 的远程机器，跳过检查")
        return 0
    failed = []
    for m in machines:
        log("CHECK", "=== 检查远程机器 {} ({}:{}, 串口 {}) ===".format(
            m.name, m.ip, m.ssh_port, m.serial_port or "-"))
        # 注意：探测命令含 shell 的 ${VAR:-default}，禁止用 str.format（花括号冲突），用拼接
        probe = (
            'echo "=USER="; whoami; '
            'echo "=NPU="; npu-smi info 2>&1 | grep -oP \'Total Count\\s*:\\s*\\K\\d+\' || npu-smi info 2>&1 | head -20; '
            'echo "=DAVINCI="; ls /dev/davinci* 2>/dev/null | grep -v davinci_manager | wc -l; '
            'echo "=ASCEND="; ls -d ${ASCEND_HOME_PATH:-/usr/local/Ascend/ascend-toolkit/latest/x86_64-linux} 2>&1; '
            'echo "=DISK="; df -h ~ | tail -1; '
            'echo "=ROOT="; mkdir -p ' + m.remote_root + ' && echo ROOT_OK; '
            'echo "=PLOG="; ls -d ' + DEFAULT_PLOG_DIR + ' 2>&1'
        )
        try:
            code, out = m.exec_cmd(probe, timeout=60, silent=True)
        except MachineError as err:
            log("CHECK", "[FAIL] {} 连接/执行失败: {}".format(m.name, err))
            failed.append(m.name)
            continue
        if code != 0:
            log("CHECK", "[FAIL] {} 探测命令退出码 {}".format(m.name, code))
            failed.append(m.name)
        for line in out.splitlines():
            if line.startswith("=") or not line.strip():
                continue
            print("    {}".format(line))
        npu_ok = re.search(r"Total Count", out) or re.search(r"Atlas|910|A5|NPU", out)
        # npu-smi 未安装时兜底：/dev/davinci* 设备数 > 0 即视为有卡（卡数取设备数）
        davinci = re.search(r"=DAVINCI=\s*\n\s*(\d+)", out)
        if not npu_ok and davinci and int(davinci.group(1)) > 0:
            npu_ok = True
            log("CHECK", "{} npu-smi 未安装，以 /dev/davinci* 设备计卡: {}".format(
                m.name, davinci.group(1)))
        ascend_ok = "ascend-toolkit" in out or "ASCEND" in out
        root_ok = "ROOT_OK" in out
        status = []
        status.append("SSH=OK")
        status.append("NPU={}".format("OK" if npu_ok else "未识别(按无卡处理,该机器不参与压测)"))
        status.append("CANN={}".format("OK" if ascend_ok else "未识别"))
        status.append("remote_root={}".format("OK" if root_ok else "FAIL"))
        log("CHECK", "{} 结果: {}".format(m.name, "; ".join(status)))
        if not root_ok:
            failed.append(m.name)
    if failed:
        log("CHECK", "[RESULT] {} 台检查异常: {}（异常机器将被剔除出压测调度）".format(
            len(failed), ",".join(sorted(set(failed)))))
        return 2
    log("CHECK", "[RESULT] 全部 {} 台远程机器检查通过".format(len(machines)))
    return 0


def cmd_exec(args):
    machines = load_machines(args.config, args.name)
    m = machines[0]
    log("EXEC", "{} ({}@{}:{}) 执行: {}".format(m.name, m.username, m.ip, m.ssh_port, args.cmd))
    code, _ = m.exec_cmd(args.cmd, timeout=args.timeout, silent=False)
    log("EXEC", "{} 退出码 {}".format(m.name, code))
    return code if code >= 0 else 1


def cmd_push(args):
    machines = load_machines(args.config, args.name)
    m = machines[0]
    remote = expand_remote_path(m, args.remote)
    m.push(args.local, remote, timeout=args.timeout)
    return 0


def cmd_pull(args):
    machines = load_machines(args.config, args.name)
    m = machines[0]
    remote = expand_remote_path(m, args.remote)
    m.pull(remote, args.local, timeout=args.timeout)
    return 0


def cmd_scan(args):
    machines = load_machines(args.config, args.name)
    plog_dir = args.plog_dir or DEFAULT_PLOG_DIR
    total_aicore = 0
    total_kernel = 0
    failed = []
    for m in machines:
        log("SCAN", "=== 远程机器 {} 扫描 (plog_dir={}) ===".format(m.name, plog_dir))
        # aicore_error 判定机制（更新版）：判定串命中，或 kernel_name + ERROR 级别报错行
        err_grep = r"\[ERROR\]|\bERROR\b"
        scans = (
            ("plog-aicore判定串", 'grep -rnE "{}" {} 2>/dev/null | head -50'.format(
                "|".join(AICORE_ERROR_PATTERNS), plog_dir), "aicore"),
            ("plog-kernel_name+ERROR", 'grep -rn "kernel_name" {} 2>/dev/null | grep -E "{}" | head -50'.format(
                plog_dir, err_grep), "aicore"),
            ("remote-output-kernel_name+ERROR",
             'grep -rn "kernel_name" {}/output 2>/dev/null | grep -E "{}" | head -50'.format(
                 m.remote_root, err_grep), "aicore"),
            ("plog-kernel_name线索(非ERROR)", 'grep -rn "kernel_name" {} 2>/dev/null | grep -vE "{}" | head -50'.format(
                plog_dir, err_grep), "kernel"),
        )
        try:
            for label, cmd, kind in scans:
                code, out = m.exec_cmd(cmd, timeout=120, silent=True)
                hits = [ln for ln in out.splitlines() if ln.strip()]
                for line in hits:
                    print("[REMOTE:{}][{}] {}".format(m.name, label, line[:300]))
                if kind == "aicore":
                    total_aicore += len(hits)
                else:
                    total_kernel += len(hits)
        except MachineError as err:
            log("SCAN", "[FAIL] {} 扫描失败: {}".format(m.name, err))
            failed.append(m.name)
    log("SCAN", "[RESULT] 远程 aicore_error命中={}（判定串或kernel_name+ERROR级别） kernel_name非ERROR线索={} 失败机器={}".format(
        total_aicore, total_kernel, ",".join(sorted(set(failed))) or "-"))
    if total_aicore > 0 or failed:
        return 2
    return 0


def main():
    parser = argparse.ArgumentParser(
        description="远程压测机器操作（list/check/exec/push/pull/scan，配置 config/remote_machines.json）")
    parser.add_argument("--config", default=os.environ.get("REMOTE_MACHINES_CONFIG", DEFAULT_CONFIG),
                        help="远程机器配置文件（默认 config/remote_machines.json，可用环境变量 REMOTE_MACHINES_CONFIG 覆盖）")
    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("list", help="列出已配置的远程机器（密码脱敏）")

    p_check = sub.add_parser("check", help="远程机器连通性与环境检查（SSH/NPU/CANN/磁盘/remote_root）")
    p_check.add_argument("--name", default=None, help="仅检查指定机器（缺省检查全部 enabled 机器）")

    p_exec = sub.add_parser("exec", help="在指定远程机器上执行命令")
    p_exec.add_argument("--name", required=True, help="机器名称")
    p_exec.add_argument("--cmd", required=True, help="要执行的远程命令（shell）")
    p_exec.add_argument("--timeout", type=int, default=600, help="超时秒数（默认 600；压测命令请配 1800+）")

    p_push = sub.add_parser("push", help="推送本地文件/目录到远程机器")
    p_push.add_argument("--name", required=True, help="机器名称")
    p_push.add_argument("--local", required=True, help="本地文件/目录")
    p_push.add_argument("--remote", required=True, help="远程目标路径（支持 ~ 与 $REMOTE_ROOT）")
    p_push.add_argument("--timeout", type=int, default=1800)

    p_pull = sub.add_parser("pull", help="拉取远程文件/目录到本地")
    p_pull.add_argument("--name", required=True, help="机器名称")
    p_pull.add_argument("--remote", required=True, help="远程文件/目录（支持 ~ 与 $REMOTE_ROOT）")
    p_pull.add_argument("--local", required=True, help="本地目标路径")
    p_pull.add_argument("--timeout", type=int, default=1800)

    p_scan = sub.add_parser("scan", help="远程 aicore_error 扫描（远程 plog + 远程 output 目录 kernel_name）")
    p_scan.add_argument("--name", default=None, help="仅扫描指定机器（缺省扫描全部 enabled 机器）")
    p_scan.add_argument("--plog_dir", default=None,
                        help="远程 plog 目录（默认 /root/ascend/log/debug/plog，可用环境变量 ASCEND_LOG_PATH 覆盖）")

    args = parser.parse_args()
    if args.command == "scan" and args.plog_dir is None:
        args.plog_dir = os.environ.get("ASCEND_LOG_PATH", DEFAULT_PLOG_DIR)

    handlers = {
        "list": cmd_list, "check": cmd_check, "exec": cmd_exec,
        "push": cmd_push, "pull": cmd_pull, "scan": cmd_scan,
    }
    try:
        return handlers[args.command](args)
    except MachineError as err:
        log("ERROR", str(err))
        return 1
    except KeyboardInterrupt:
        log("ERROR", "收到中断退出")
        return 1


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