"""本地测试用 SSH 服务端:仅用于自测隧道引擎。
监听 127.0.0.1:2222,接受用户 test / 密码 test123,支持 direct-tcpip 端口转发。
用法: python test_sshd.py
"""
import select
import socket
import threading
from pathlib import Path

import paramiko

HOST, PORT = "127.0.0.1", 2222
USER, PW = "test", "test123"
KEY_PATH = Path(__file__).parent / ".test_host_key"


def get_host_key() -> paramiko.RSAKey:
    if KEY_PATH.exists():
        return paramiko.RSAKey.from_private_key_file(str(KEY_PATH))
    key = paramiko.RSAKey.generate(2048)
    key.write_private_key_file(str(KEY_PATH))
    return key


class Server(paramiko.ServerInterface):
    def __init__(self):
        self.dest: dict[int, tuple] = {}

    def get_allowed_auths(self, username):
        return "password"

    def check_auth_password(self, username, password):
        return paramiko.AUTH_SUCCESSFUL if (username, password) == (USER, PW) else paramiko.AUTH_FAILED

    def check_channel_request(self, kind, chanid):
        return paramiko.OPEN_SUCCEEDED

    def check_channel_direct_tcpip_request(self, chanid, origin, destination):
        self.dest[chanid] = tuple(destination)
        return paramiko.OPEN_SUCCEEDED


def pump(a, b) -> None:
    try:
        while True:
            r, _, _ = select.select([a, b], [], [], 1.0)
            if not r:
                if hasattr(a, "closed") and a.closed:
                    break
                continue
            if b in r:
                data = b.recv(65536)
                if not data:
                    break
                a.sendall(data)
            if a in r:
                data = a.recv(65536)
                if not data:
                    break
                b.sendall(data)
    except Exception:
        pass
    finally:
        try: a.close()
        except Exception: pass
        try: b.close()
        except Exception: pass


def handle(conn: socket.socket, key) -> None:
    t = paramiko.Transport(conn)
    t.add_server_key(key)
    server = Server()
    try:
        t.start_server(server=server)
    except paramiko.SSHException:
        return
    chan = t.accept(20)
    if chan is None:
        return
    dest = server.dest.get(chan.get_id())
    if not dest:
        chan.close()
        return
    try:
        upstream = socket.create_connection(dest, timeout=10)
    except OSError:
        chan.close()
        return
    pump(chan, upstream)


def main() -> None:
    key = get_host_key()
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((HOST, PORT))
    srv.listen(8)
    print(f"test sshd listening on {HOST}:{PORT} (login {USER}/{PW})")
    try:
        while True:
            conn, _addr = srv.accept()
            threading.Thread(target=handle, args=(conn, key), daemon=True).start()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()