github copilot生成的paramiko shell方法

#!/usr/bin/env python3
import paramiko
import socket
import time
import uuid

# ------------ 硬编码登录信息（请按需修改） ------------
HOST = "example.com"
PORT = 22
USERNAME = "youruser"
PASSWORD = "yourpassword"
# -----------------------------------------------------

class SSHPersistentShell:
    """
    长连接 + 持久交互 shell 的轻量客户端。
    用法：
      c = SSHPersistentShell()
      c.connect()
      out, exit_code = c.send_and_receive("ls -l /tmp", timeout=5.0)
      c.close()
    注意：该类非线程安全；如果并发调用 send_and_receive，需要外部加锁。
    """
    def __init__(self, host=HOST, port=PORT, username=USERNAME, password=PASSWORD,
                 connect_timeout=5.0, keepalive=30, poll_interval=0.002):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.connect_timeout = connect_timeout
        self.keepalive = keepalive
        self.poll_interval = poll_interval  # 轮询间隔（秒），越小延迟越低但 CPU 占用更高

        self._sock = None
        self._transport = None
        self._chan = None

    def connect(self):
        """建立 Transport 并 invoke_shell() 建立持久 shell channel。"""
        if self.is_connected():
            return

        sock = None
        transport = None
        try:
            sock = socket.create_connection((self.host, self.port), timeout=self.connect_timeout)
            sock.settimeout(None)
            transport = paramiko.Transport(sock)
            transport.start_client(timeout=self.connect_timeout)
            transport.auth_password(self.username, self.password, timeout=self.connect_timeout)

            # 打开持久 shell
            chan = transport.open_session()
            # 请求 pty 以获得交互式 shell 行为（许多命令需要 tty）
            chan.get_pty(term='vt100', width=80, height=24)
            chan.invoke_shell()
            chan.settimeout(0.0)  # non-blocking recv

            # 基本环境初始化：统一 locale，清除 PROMPT_COMMAND、把 PS1 设空、关闭 echo，减少噪音和回显
            init_cmds = (
                "export LC_ALL=C >/dev/null 2>&1 || true; "
                "unset PROMPT_COMMAND >/dev/null 2>&1 || true; "
                "PS1='' >/dev/null 2>&1 || true; "
                "stty -echo >/dev/null 2>&1 || true\n"
            )
            chan.send(init_cmds)

            # Drain initial data (login banner / prompt) — 读取短时间直到没有新数据
            buf = bytearray()
            t0 = time.time()
            last_read = time.time()
            while time.time() - t0 < 2.0:  # 最多等 2 秒用于清空初始输出
                if chan.recv_ready():
                    data = chan.recv(65536)
                    if not data:
                        break
                    buf.extend(data)
                    last_read = time.time()
                    # 继续读取，直到 0.05s 没有新数据
                else:
                    if time.time() - last_read > 0.05:
                        break
                    time.sleep(0.01)

            # 保存资源
            self._sock = sock
            self._transport = transport
            self._chan = chan

            try:
                transport.set_keepalive(self.keepalive)
            except Exception:
                pass

        except Exception:
            # 清理资源
            if transport is not None:
                try:
                    transport.close()
                except Exception:
                    pass
            if sock is not None:
                try:
                    sock.close()
                except Exception:
                    pass
            raise

    def is_connected(self):
        return (self._transport is not None) and (self._transport.is_active()) and (self._chan is not None) and (not self._chan.closed)

    def send_and_receive(self, command: str, timeout: float = 10.0):
        """
        在持久 shell 上发送命令并等待完成。
        返回 (combined_output_str, exit_status:int).
        timeout 秒后抛出 TimeoutError（不会自动断开上层连接，但会尝试关闭当前 channel 并让用户 decide）。
        PS: stdout/stderr 在 pty 模式下会合并到一起。
        """
        if not self.is_connected():
            raise RuntimeError("Not connected; call connect() first.")

        # 唯一结束标记，确保在输出中唯一
        marker = uuid.uuid4().hex
        sentinel = f"__PYSSH_DONE_{marker}__"
        sentinel_b = sentinel.encode()

        # 在命令后发送 printf 来打印 sentinel 和 exit code（以便我们检测命令何时结束）
        # 注意：2>&1 没必要，因为 pty 已合并 stdout/stderr
        payload = f"{command}\nprintf '\\n{sentinel}:%s\\n' $? \n"
        try:
            self._chan.send(payload)
        except Exception as e:
            raise RuntimeError("Failed to send command to channel: " + str(e))

        outBuf = bytearray()
        start = time.time()

        # 读取直到看到 sentinel
        while True:
            # 先读取所有可用数据以降低延迟
            try:
                if self._chan.recv_ready():
                    data = self._chan.recv(65536)
                    if data:
                        outBuf.extend(data)
                # 检查 sentinel 是否出现
                idx = outBuf.find(sentinel_b)
                if idx != -1:
                    # sentinel 后面应该是 b':<exitcode>\\n'
                    after = outBuf[idx + len(sentinel_b):]
                    # 查找冒号 (should be ':') 和换行
                    # 我们 expect pattern: sentinel:NNN\n
                    # 寻找换行位置
                    nl = after.find(b'\n')
                    if nl == -1:
                        # 还没完整接收到整行，继续读
                        pass
                    else:
                        # 提取 exit code 字段（在冒号后到换行）
                        # after can start with b':', so skip leading b':'
                        code_field = after
                        if code_field.startswith(b':'):
                            code_field = code_field[1:]
                        # code_field may be like b'0\r' etc; take up to nl
                        code_str = code_field[:nl].strip().decode(errors='ignore')
                        try:
                            exit_code = int(code_str) if code_str else 0
                        except Exception:
                            exit_code = -1
                        # 输出内容是 outBuf up to idx (剔除 sentinel 行及之后的内容)
                        output_bytes = outBuf[:idx]
                        output = output_bytes.decode(errors='replace')
                        return output, exit_code

                # timeout check
                if (time.time() - start) > timeout:
                    # 不直接关闭整个 transport，但尝试关闭 channel 并 raise
                    try:
                        self._chan.close()
                    except Exception:
                        pass
                    raise TimeoutError(f"Command timeout after {timeout:.3f}s: {command!r}")

                # 短等待避免 busy loop
                time.sleep(self.poll_interval)
            except Exception as e:
                # 遇到 socket/transport 层面异常，重新抛出
                raise RuntimeError("Error while receiving: " + str(e))

    def close(self):
        """关闭 channel/transport/socket。"""
        if self._chan is not None:
            try:
                self._chan.close()
            except Exception:
                pass
            self._chan = None
        if self._transport is not None:
            try:
                self._transport.close()
            except Exception:
                pass
            self._transport = None
        if self._sock is not None:
            try:
                self._sock.close()
            except Exception:
                pass
            self._sock = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.close()


if __name__ == "__main__":
    # 示例：把 HOST/USERNAME/PASSWORD 改成真实值后运行
    client = SSHPersistentShell()
    try:
        client.connect()
        # 多次短命令示例，见延迟很小
        for cmd in ["echo hello", "date +%s.%N", "ls -1 /tmp | wc -l"]:
            out, code = client.send_and_receive(cmd, timeout=5.0)
            print(f"CMD: {cmd!r} -> exit:{code}\n{out}")
    finally:
        client.close()