ykxmz/_build/lssh.py

# -*- coding: utf-8 -*-
r'''连 AWS Lightsail 那台新服务器的 SSH 助手。

== 为什么单独一个模块，不改 devssh.py ==

devssh.py 是连**随身 WiFi 那台 4G 小机器**的（192.168.2.156，密码登录）。
那台还在跑（网站暂时还挂在上面），不能动它的连接逻辑。

新服务器用**密钥**登录、IP 是公网 <源站IP>，认证方式完全不同，
所以单独一个模块更清楚。

== 用法 ==

    from lssh import run, upload, download, probe
    run("uname -a")                     # 跑命令
    upload("relay/relay.py", "/opt/ykxmz-relay/relay.py")
'''
from __future__ import annotations

import os
import socket
from pathlib import Path

import paramiko

HOST = os.environ.get("YKXMZ_LS_HOST", "<源站IP>")
USER = os.environ.get("YKXMZ_LS_USER", "admin")
PORT = int(os.environ.get("YKXMZ_LS_PORT", "22"))
KEY = Path(os.environ.get(
    "YKXMZ_LS_KEY", str(Path.home() / ".ssh" / "lightsail-ykxmz.pem")))

_client = None


def _connect(timeout: int = 25):
    global _client
    if _client is not None:
        try:
            # 探一下还活着没
            _client.exec_command("true", timeout=5)
            return _client
        except Exception:                                 # noqa: BLE001
            _client = None

    if not KEY.is_file():
        raise RuntimeError(
            "找不到 SSH 密钥：%s\n"
            "（把 Lightsail 的 .pem 放到那里，或者设 YKXMZ_LS_KEY）" % KEY)

    cli = paramiko.SSHClient()
    cli.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    cli.connect(HOST, port=PORT, username=USER,
                key_filename=str(KEY), timeout=timeout,
                allow_agent=False, look_for_keys=False,
                banner_timeout=30, auth_timeout=30)
    _client = cli
    return cli


def run(command: str, timeout: int = 300, quiet: bool = False) -> tuple:
    """跑一条命令，返回 (退出码, 输出)。"""
    cli = _connect()
    if not quiet:
        print("  [ls] %s" % command.splitlines()[0][:90])
    stdin, stdout, stderr = cli.exec_command(command, timeout=timeout,
                                             get_pty=False)
    out = stdout.read().decode("utf-8", "replace")
    err = stderr.read().decode("utf-8", "replace")
    code = stdout.channel.recv_exit_status()
    return code, (out + err)


def upload(local: str, remote: str, mode: int = 0o644) -> None:
    """传一个文件上去。"""
    cli = _connect()
    sftp = cli.open_sftp()
    try:
        # 先把远端目录建出来
        d = os.path.dirname(remote)
        if d:
            try:
                sftp.stat(d)
            except IOError:
                # 逐级建
                parts = d.strip("/").split("/")
                cur = ""
                for p in parts:
                    cur += "/" + p
                    try:
                        sftp.stat(cur)
                    except IOError:
                        try:
                            sftp.mkdir(cur)
                        except IOError:
                            pass
        sftp.put(local, remote)
        sftp.chmod(remote, mode)
    finally:
        sftp.close()
    print("  已上传 %s -> %s" % (os.path.basename(local), remote))


def download(remote: str, local: str) -> None:
    cli = _connect()
    sftp = cli.open_sftp()
    try:
        Path(local).parent.mkdir(parents=True, exist_ok=True)
        sftp.get(remote, local)
    finally:
        sftp.close()
    print("  已下载 %s -> %s" % (remote, local))


def probe() -> str | None:
    """通不通。返回一句话描述，或 None。"""
    try:
        code, out = run("uname -sm; nproc; free -m | sed -n 2p",
                        timeout=30, quiet=True)
        if code != 0:
            return None
        return out.strip().replace("\n", " | ")
    except Exception as exc:                              # noqa: BLE001
        return "连接失败：%s" % exc


if __name__ == "__main__":
    print("  目标: %s@%s:%d  密钥 %s" % (USER, HOST, PORT, KEY))
    print("  " + (probe() or "不通"))