ssh-e2e/test_e2e.py

"""端到端验证:走真实的 Api 代码路径建立隧道,再请求经隧道转发的 HTTP 服务。
前置: 先运行 `python test_sshd.py`(127.0.0.1:2222, test/test123)。
用法: .venv/Scripts/python test_e2e.py
"""
import socket
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).parent
sys.path.insert(0, str(ROOT))

TEST_DIR = ROOT / ".test_data"
LOCAL_PORT = 13000   # 隧道本机端口
SVC_PORT = 3210      # 模拟"只对远程主机本机开放"的服务端口


def wait_port(port: int, expect_open: bool, timeout: float = 15.0) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        s = socket.socket()
        s.settimeout(0.4)
        try:
            s.connect(("127.0.0.1", port))
            s.close()
            if expect_open:
                return True
        except OSError:
            if not expect_open:
                return True
        finally:
            s.close()
        time.sleep(0.25)
    return False


def tunnel_status(api, tid: str) -> str:
    for t in api.get_state()["state"]["tunnels"]:
        if t["id"] == tid:
            return t["runtime"]["status"]
    return "gone"


def main() -> int:
    import shutil
    shutil.rmtree(TEST_DIR, ignore_errors=True)
    TEST_DIR.mkdir(parents=True)

    from storage import Store
    from engine import TunnelManager
    from main import Api

    events = []
    manager = TunnelManager(lambda ev: events.append(ev), TEST_DIR)
    store = Store(TEST_DIR)
    api = Api(store, manager)

    # 1) 后台起"远程服务"
    www = TEST_DIR / "www"
    www.mkdir(exist_ok=True)
    (www / "index.html").write_text("<h1>TUNNEL-E2E-OK</h1>", encoding="utf-8")
    svc = subprocess.Popen(
        [sys.executable, "-m", "http.server", str(SVC_PORT), "--directory", str(www)],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    assert wait_port(SVC_PORT, True), "http 服务未启动"

    # 2) 起本地测试 SSH 服务端
    sshd = subprocess.Popen([sys.executable, str(ROOT / "test_sshd.py")],
                            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    assert wait_port(2222, True), "test_sshd 未启动"

    try:
        # 3) 通过 Api 添加主机和隧道
        r = api.save_host({"name": "本机测试", "hostname": "127.0.0.1", "port": 2222,
                           "username": "test", "auth": "password", "password": "test123"})
        assert r["ok"], r
        host_id = r["id"]

        # 密码错误的情况应失败
        r = api.probe_host({"hostname": "127.0.0.1", "port": 2222, "username": "test",
                            "auth": "password", "password": "wrong"})
        assert r["ok"] is False, f"错误密码竟能通过: {r}"
        print("[PASS] 错误密码被正确拒绝:", r["message"])

        r = api.probe_host({"hostname": "127.0.0.1", "port": 2222, "username": "test",
                            "auth": "password", "password": "test123"})
        assert r["ok"], r
        print("[PASS] 主机测试连接:", r["message"])

        r = api.save_tunnel({"host_id": host_id, "name": "e2e", "local_port": LOCAL_PORT,
                             "remote_host": "127.0.0.1", "remote_port": SVC_PORT,
                             "autostart": False, "autoreconnect": True})
        assert r["ok"], r
        tid = r["id"]

        # 4) 启动隧道并等待建立
        assert api.start_tunnel(tid)["ok"]
        deadline = time.time() + 20
        while time.time() < deadline and tunnel_status(api, tid) != "connected":
            time.sleep(0.3)
        assert tunnel_status(api, tid) == "connected", \
            "隧道未能建立: " + json_events_tail(events)
        print(f"[PASS] 隧道已建立 127.0.0.1:{LOCAL_PORT} → 127.0.0.1:{SVC_PORT}")

        # 5) 经隧道请求 HTTP 服务
        with urllib.request.urlopen(f"http://127.0.0.1:{LOCAL_PORT}", timeout=10) as resp:
            body = resp.read().decode()
        assert "TUNNEL-E2E-OK" in body, body
        print("[PASS] 经隧道取回内容:", body.strip())

        # 6) 断开后端口应关闭
        api.stop_tunnel(tid)
        deadline = time.time() + 10
        while time.time() < deadline and tunnel_status(api, tid) != "stopped":
            time.sleep(0.2)
        assert tunnel_status(api, tid) == "stopped"
        assert wait_port(LOCAL_PORT, False, timeout=5), "断开后本地端口仍可连接"
        print("[PASS] 断开隧道,本地端口已释放")

        # 7) DPAPI 加密落盘校验
        store.save()
        raw = (TEST_DIR / "config.json").read_text(encoding="utf-8")
        assert "test123" not in raw and "dpapi:" in raw
        print("[PASS] 密码已以 DPAPI 加密存储,明文不落盘")

        print("\n全部通过 ✔")
        return 0
    finally:
        manager.stop_all()
        svc.terminate()
        sshd.terminate()


def json_events_tail(events) -> str:
    import json
    lines = [e for e in events if e.get("type") == "log"][-6:]
    return "\n".join(json.dumps(e, ensure_ascii=False) for e in lines)


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