ykxmz/_build/diag_relay_leg.py

# -*- coding: utf-8 -*-
"""单独验证"跨网络中转载荷"这条腿，每步都打时间戳，定位卡在哪。

只读诊断：起中转 + 真实被控端，然后用 HTTP 长轮询走一遍 pair/ping。
"""
from __future__ import annotations

import json
import os
import shutil
import socket
import subprocess
import sys
import tempfile
import time
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "windows"))

_UA = {"User-Agent": "LanRemote/1.8.2"}
T0 = time.time()


def say(msg: str) -> None:
    print("  [%6.1fs] %s" % (time.time() - T0, msg), flush=True)


def free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def post(base, path, body=b"", headers=None, timeout=20.0):
    req = urllib.request.Request(base + path, data=body, method="POST",
                                 headers=dict(_UA))
    for k, v in (headers or {}).items():
        req.add_header(k, v)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.status, resp.read()


def main() -> int:
    work = Path(tempfile.mkdtemp(prefix="diag_relayleg_"))
    relay = agent = None
    try:
        from lanremote.config import Store
        relay_port = free_port()
        agent_port = free_port()
        base = "http://127.0.0.1:%d" % relay_port
        code = "246813"

        say("起中转 %s" % base)
        rd = work / "relaydata"
        rd.mkdir(parents=True, exist_ok=True)
        env = dict(os.environ)
        env.update({"RELAY_PORT": str(relay_port), "RELAY_BIND": "127.0.0.1",
                    "RELAY_DATA": str(rd), "RELAY_POLL": "2",
                    "PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"})
        relay = subprocess.Popen(
            [sys.executable, "-u", str(ROOT / "relay" / "relay.py")], env=env,
            stdout=open(work / "relay.log", "wb"), stderr=subprocess.STDOUT)
        for _ in range(60):
            time.sleep(0.5)
            try:
                with urllib.request.urlopen(base + "/v1/health", timeout=2):
                    break
            except Exception:                               # noqa: BLE001
                continue
        say("中转就绪")

        data = work / "agentdata"
        data.mkdir(parents=True, exist_ok=True)
        st = Store(data)
        st.set_pairing_code(code)
        st.set("setup_done", True)
        st.set("allow_new_pairing", True)
        st.set("close_pair_after_first", False)
        st.set("server_name", "诊断电脑")
        st.set("relay_url", base)
        st.set("port", agent_port)

        env2 = dict(os.environ)
        env2.update({"YKXMZ_DATA": str(data),
                     "YKXMZ_LOCK_PORT": str(free_port()),
                     "PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"})
        say("起被控端 端口 %d" % agent_port)
        agent = subprocess.Popen(
            [sys.executable, "-u", str(ROOT / "windows" / "run_agent.py"),
             "--silent", "--port", str(agent_port)],
            cwd=str(ROOT / "windows"), env=env2,
            stdout=open(work / "agent.log", "wb"), stderr=subprocess.STDOUT)

        online = False
        for _ in range(45):
            time.sleep(1.0)
            try:
                with urllib.request.urlopen(base + "/v1/health", timeout=4) as r:
                    snap = json.loads(r.read().decode("utf-8"))
            except Exception:                               # noqa: BLE001
                continue
            if any(d.get("online") for d in snap.get("devices", [])):
                online = True
                break
        say("被控端在线: %s" % online)
        if not online:
            say((data / "silent.log").read_text(encoding="utf-8",
                                                errors="replace")[-500:])
            return 1

        say("attach 接入")
        body = json.dumps({"code": code, "name": "诊断手表"},
                          ensure_ascii=False).encode()
        _s, raw = post(base, "/v1/client/attach", body)
        attach = json.loads(raw.decode())
        session = attach.get("session", "")
        say("attach ok=%s session=%s" % (attach.get("ok"), session[:12]))

        buf = b""

        def call(obj, timeout=25.0):
            nonlocal buf
            payload = (json.dumps(obj, ensure_ascii=False) + "\n").encode()
            say("  send %s" % obj.get("cmd"))
            post(base, "/v1/client/send", payload, {"X-Session": session},
                 timeout)
            say("  send 返回，开始 poll")
            deadline = time.time() + timeout
            while True:
                while b"\n" in buf:
                    line, _, buf = buf.partition(b"\n")
                    if not line.strip():
                        continue
                    msg = json.loads(line.decode())
                    if "event" in msg:
                        continue
                    say("  收到 %s" % msg.get("cmd"))
                    return msg
                left = deadline - time.time()
                if left <= 0:
                    raise TimeoutError("等回复超时")
                t = time.time()
                _s, chunk = post(base, "/v1/client/poll", b"",
                                 {"X-Session": session}, min(20.0, left))
                say("  poll 返回 %d 字节，用了 %.1fs"
                    % (len(chunk), time.time() - t))
                buf += chunk

        say("pair")
        reply = call({"cmd": "pair", "deviceId": "diag-relay",
                      "deviceName": "诊断手表", "code": code})
        say("pair ok=%s" % reply.get("ok"))

        for i in range(3):
            t = time.time()
            pong = call({"cmd": "ping"})
            say("ping #%d ok=%s 往返 %.0f ms"
                % (i + 1, pong.get("ok"), (time.time() - t) * 1000))
        return 0
    except Exception as exc:                                # noqa: BLE001
        import traceback
        say("异常: %s" % exc)
        traceback.print_exc()
        return 1
    finally:
        for p in (agent, relay):
            if p is not None:
                try:
                    p.terminate()
                    p.wait(timeout=5)
                except Exception:                           # noqa: BLE001
                    try:
                        p.kill()
                    except Exception:                      # noqa: BLE001
                        pass
        time.sleep(0.3)
        shutil.rmtree(work, ignore_errors=True)


if __name__ == "__main__":
    raise SystemExit(main())