ykxmz/_build/ins_relay_flow.py

# -*- coding: utf-8 -*-
r'''给中转打临时插桩，抓「数据到底去哪了」。

做法：把 relay.py 复制一份到临时目录，在关键位置插入 log 调用：
  * /v1/client/send   -> 谁发的、多少字节、当时 to_agent 有多长
  * /v1/agent/poll    -> 取出多少字节、当时的 X-Relay-Session 是什么
  * /v1/agent/send    -> 被控端回了多少字节
  * /v1/client/poll   -> 手表取走多少字节
  * _sync_session     -> 会话令牌怎么变的

然后在同一个 repro 里跑，日志能直接指出哪一步数据消失了。
'''
from __future__ import annotations

import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
WINDOWS = ROOT / "windows"
RELAY_PORT = 8126
AGENT_PORT = 8127
BASE = f"http://127.0.0.1:{RELAY_PORT}"
PAIR = "135790"
UA = {"User-Agent": "LanRemote/1.8.2"}


def instrument(src: str) -> str:
    """在 relay.py 的关键位置插 log。"""
    # 1. client/send
    src = src.replace(
        """            data = self._body()
            RELAY.client_put(session, data)
            self._json(200, {"ok": True, "queued": len(data)})""",
        """            data = self._body()
            _dev = RELAY.devices.get(session.device_id)
            log(f"[INS] client/send {len(data)}B -> to_agent 现有 "
                f"{len(_dev.to_agent) if _dev else '?'}B，会话 {token[:8]}，"
                f"current_session={( _dev.current_session or '-')[:8] if _dev else '?'}")
            RELAY.client_put(session, data)
            self._json(200, {"ok": True, "queued": len(data)})""", 1)

    # 2. agent/poll
    src = src.replace(
        """            session_at_start = device.current_session
            data = RELAY.agent_take(device, POLL_SECONDS)""",
        """            session_at_start = device.current_session
            log(f"[INS] agent/poll 进入，to_agent={len(device.to_agent)}B，"
                f"current_session={(device.current_session or '-')[:8]}")
            data = RELAY.agent_take(device, POLL_SECONDS)
            log(f"[INS] agent/poll 返回 {len(data)}B，"
                f"current_session={(device.current_session or '-')[:8]}")""", 1)

    # 3. agent/send
    src = src.replace(
        """            data = self._body()
            RELAY.agent_put(device, data)
            self._json(200, {"ok": True, "queued": len(data)})""",
        """            data = self._body()
            log(f"[INS] agent/send {len(data)}B，active_session="
                f"{(device.current_session or '-')[:8]}")
            RELAY.agent_put(device, data)
            self._json(200, {"ok": True, "queued": len(data)})""", 1)

    # 4. client/poll
    src = src.replace(
        """from __future__ import annotations""",
        """from __future__ import annotations  # instrumented""", 1)
    return src


def post(path, body=b"", headers=None, timeout=20.0):
    h = dict(UA)
    h.update(headers or {})
    req = urllib.request.Request(BASE + path, data=body, headers=h, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read()
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read()
    except Exception as exc:                              # noqa: BLE001
        return 0, str(exc).encode()


def get(path, timeout=10.0):
    req = urllib.request.Request(BASE + path, headers=UA)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read()
    except Exception as exc:                              # noqa: BLE001
        return 0, str(exc).encode()


class Watch:
    def __init__(self, session):
        self.session = session
        self.buf = b""

    def send(self, obj):
        payload = (json.dumps(obj, ensure_ascii=False) + "\n").encode()
        st, _ = post("/v1/client/send", payload, {"X-Session": self.session})
        return st == 200

    def reply(self, timeout=20.0):
        deadline = time.time() + timeout
        while time.time() < deadline:
            while b"\n" in self.buf:
                line, _, self.buf = self.buf.partition(b"\n")
                if not line.strip():
                    continue
                msg = json.loads(line.decode("utf-8"))
                if "event" in msg:
                    continue
                return msg
            st, body = post("/v1/client/poll", b"", {"X-Session": self.session},
                            timeout=max(4.0, deadline - time.time()))
            if st != 200:
                time.sleep(0.4)
                continue
            self.buf += body
        return None


def main() -> int:
    work = Path(tempfile.mkdtemp(prefix="ins_"))
    relay_proc = agent_proc = None
    real_data = WINDOWS / "data"
    backup = WINDOWS / "data_bk_ins"
    had = real_data.exists()
    if had:
        if backup.exists():
            shutil.rmtree(backup, ignore_errors=True)
        shutil.move(str(real_data), str(backup))
    try:
        # 插桩版的 relay.py
        relay_src = (ROOT / "relay" / "relay.py").read_text(encoding="utf-8")
        inst = instrument(relay_src)
        if "[INS]" not in inst:
            print("  ✗ 插桩失败：锚点没匹配上")
            return 1
        n = inst.count("[INS]")
        print("  插桩了 %d 处" % n)
        relay_dir = work / "relaydir"
        shutil.copytree(ROOT / "relay", relay_dir)
        (relay_dir / "relay.py").write_text(inst, encoding="utf-8",
                                            newline="\n")

        real_data.mkdir(parents=True, exist_ok=True)
        sys.path.insert(0, str(WINDOWS))
        from lanremote.config import Store
        store = Store(real_data)
        store.set_pairing_code(PAIR)
        store.set("setup_done", True)
        store.set("allow_new_pairing", True)
        store.set("close_pair_after_first", False)
        store.set("server_name", "插桩电脑")
        store.set("relay_url", f"http://127.0.0.1:{RELAY_PORT}")
        store.set("port", AGENT_PORT)

        rdata = work / "r"
        rdata.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(rdata), "RELAY_POLL": "5",
                    "PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"})
        # ⚠ 用户自己的被控端正跑在这台机器上，占着单实例锁端口 8767。
        #   不换端口的话测试实例一起来就被"已有实例在运行"顶掉，
        #   表现为"被控端根本没登记"，而且日志是空的，极难排查。
        agent_env = dict(env)
        agent_env["YKXMZ_LOCK_PORT"] = "18767"
        relay_log = open(work / "relay.log", "wb")
        relay_proc = subprocess.Popen(
            [sys.executable, "-u", str(relay_dir / "relay.py")], env=env,
            cwd=str(relay_dir), stdout=relay_log, stderr=subprocess.STDOUT)
        for _ in range(40):
            time.sleep(0.25)
            if get("/v1/health")[0] == 200:
                break
        print("  插桩中转起来了")

        agent_log = open(work / "agent.log", "wb")
        agent_proc = subprocess.Popen(
            [sys.executable, "-u", str(WINDOWS / "run_agent.py"),
             "--silent", "--port", str(AGENT_PORT)],
            cwd=str(WINDOWS), stdout=agent_log, stderr=subprocess.STDOUT,
            env=agent_env)

        device = None
        for _ in range(40):
            time.sleep(0.5)
            st, body = get("/v1/health")
            if st == 200:
                snap = json.loads(body.decode())
                if snap.get("devices"):
                    device = snap["devices"][0]
                    break
        print("  被控端登记: %s" % (device["name"] if device else "**没有**"))

        st, body = post("/v1/client/attach", json.dumps(
            {"code": PAIR, "name": "插桩手表"}).encode(),
            {"Content-Type": "application/json"})
        att = json.loads(body.decode())
        watch = Watch(att.get("session", ""))
        print("  手表接入: ok=%s" % att.get("ok"))

        steps = [{"cmd": "pair", "deviceId": "w1", "name": "手表", "code": PAIR},
                 {"cmd": "ping"}, {"cmd": "ping"}, {"cmd": "ping"}]
        for i, msg in enumerate(steps, 1):
            t0 = time.time()
            watch.send(msg)
            got = watch.reply(timeout=15.0)
            dt = int((time.time() - t0) * 1000)
            print("  %d. %-8s %s  %dms" % (i, msg["cmd"],
                  "OK" if got else "**超时**", dt))
            if got is None:
                print("     卡住了，停止")
                break
        time.sleep(1.0)

        print()
        print("  ════ 插桩日志（中转侧数据流）════")
        relay_log.flush()
        txt = (work / "relay.log").read_text(encoding="utf-8", errors="replace")
        for line in txt.splitlines():
            if "[INS]" in line:
                print("    " + line[:130])

        print()
        print("  ════ 被控端日志 ════")
        agent_log.flush()
        txt = (work / "agent.log").read_text(encoding="utf-8", errors="replace")
        for line in txt.splitlines()[-20:]:
            print("    " + line[:120])
        return 0
    finally:
        for p in (agent_proc, relay_proc):
            if p is not None:
                try:
                    p.terminate()
                    p.wait(timeout=6)
                except Exception:                         # noqa: BLE001
                    try:
                        p.kill()
                    except Exception:                     # noqa: BLE001
                        pass
        shutil.rmtree(work, ignore_errors=True)
        shutil.rmtree(real_data, ignore_errors=True)
        if had:
            shutil.move(str(backup), str(real_data))


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