ykxmz/_build/devssh.py

# -*- coding: utf-8 -*-
"""随身 WiFi（Debian）远程执行小工具。

用法：
    python _build/devssh.py "命令"            # 跑一条命令
    python _build/devssh.py --script 文件     # 跑一段脚本

== 地址会自动试 ==

这台设备跟着教室网络走，地址会变：
    192.168.2.156  <->  192.168.10.105
所以这里**按顺序试多个地址**，谁先连通就用谁，不用改代码。

想强制指定就用环境变量：
    $env:YKXMZ_DEV_HOST = "192.168.10.105"

⚠ 以前的坑：定义了 HOST_FALLBACKS 但**三个函数都直接用 HOST**，
  那个常量根本没被用过。于是设备一换网段就全部连不上，
  而且报的是「Connection refused / timed out」，看起来像设备挂了。
"""
from __future__ import annotations

import os
import sys
import time
from pathlib import Path

import paramiko

USER = "root"
PASSWORD = os.environ.get("YKXMZ_DEV_PASS", "1313144")

# 候选地址，按顺序试。
#   环境变量优先（想锁定时用），然后是最近见过的两个地址。
_CANDIDATES: list[str] = []


def _build_candidates() -> list[str]:
    out = []
    env_host = (os.environ.get("YKXMZ_DEV_HOST") or "").strip()
    if env_host:
        out.append(env_host)
    for host in ("192.168.2.156",      # 当前地址
                 "192.168.10.105"):    # 换网络前用过
        if host not in out:
            out.append(host)
    return out


CANDIDATES = _build_candidates()
HOST = CANDIDATES[0]          # 兼容老代码：默认取第一个
HOST_FALLBACKS = tuple(CANDIDATES[1:])

# 记住这次连通的是哪个，后续调用不用每次从头试
_CONNECTED: str | None = None


def _connect(timeout: int = 15):
    """按候选顺序连，返回 (client, host)。全都不通才抛异常。"""
    global _CONNECTED
    order = list(CANDIDATES)
    if _CONNECTED and _CONNECTED in order:
        # 上次通的那个排最前，省一次失败尝试
        order.remove(_CONNECTED)
        order.insert(0, _CONNECTED)

    errors = []
    for host in order:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            client.connect(host, username=USER, password=PASSWORD,
                           timeout=timeout, banner_timeout=20,
                           auth_timeout=20, look_for_keys=False,
                           allow_agent=False)
            if _CONNECTED != host:
                print(f"  [devssh] 连上 {host}")
            _CONNECTED = host
            return client, host
        except Exception as exc:                          # noqa: BLE001
            errors.append(f"{host}: {type(exc).__name__} {exc}")
            try:
                client.close()
            except Exception:                             # noqa: BLE001
                pass
    raise RuntimeError("所有候选地址都连不上：\n  " + "\n  ".join(errors))


def probe() -> str | None:
    """探一下现在哪个地址能用（不抛异常）。"""
    try:
        client, host = _connect(timeout=8)
        client.close()
        return host
    except Exception:                                     # noqa: BLE001
        return None


def run(command: str, timeout: int = 600, quiet: bool = False) -> tuple:
    client, _host = _connect()
    try:
        _stdin, stdout, stderr = client.exec_command(command, timeout=timeout)
        out = stdout.read().decode("utf-8", "replace")
        err = stderr.read().decode("utf-8", "replace")
        code = stdout.channel.recv_exit_status()
        if not quiet:
            if out.strip():
                print(out.rstrip())
            if err.strip():
                print("[stderr] " + err.rstrip())
        return code, out, err
    finally:
        client.close()


def upload(local: str, remote: str, mode: int = 0o644) -> None:
    client, _host = _connect()
    try:
        sftp = client.open_sftp()
        sftp.put(local, remote)
        sftp.chmod(remote, mode)
        sftp.close()
        print(f"  已上传 {Path(local).name} -> {remote}")
    finally:
        client.close()


def download(remote: str, local: str) -> None:
    client, _host = _connect()
    try:
        sftp = client.open_sftp()
        sftp.get(remote, local)
        sftp.close()
        print(f"  已下载 {remote} -> {local}")
    finally:
        client.close()


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        print("  候选地址: %s" % ", ".join(CANDIDATES))
        print("  现在能连的: %s" % (probe() or "（一个都不通）"))
        sys.exit(1)
    if sys.argv[1] == "--script":
        run(Path(sys.argv[2]).read_text(encoding="utf-8"), timeout=600)
    elif sys.argv[1] == "--probe":
        got = probe()
        print("  能连的地址: %s" % (got or "（一个都不通）"))
        sys.exit(0 if got else 1)
    else:
        args = sys.argv[1:]
        timeout = 600
        if args and args[0] == "--timeout":
            timeout = int(args[1])
            args = args[2:]
        sys.exit(run(" ".join(args), timeout=timeout)[0])