ykxmz/windows/lanremote/discovery.py

# -*- coding: utf-8 -*-
"""局域网自动发现（UDP 广播）。

目的：**手表端不用输 IP**。
手表往局域网广播一个探测包，被控端收到后单播回自己的名字和端口，
手表就能自动填好地址直接连。

这样也顺手解决了另一个坑：电脑换了网络、IP 变了、或者机器上有
Tailscale / VMware 之类虚拟网卡时，用户猜不准该填哪个 IP。
"""
from __future__ import annotations

import json
import socket

from . import __version__ as _APP_VERSION
import threading
import time

DISCOVERY_PORT = 8766
PROBE_MAGIC = "LANREMOTE_DISCOVER"
REPLY_MAGIC = "LANREMOTE_REPLY"
# 手表点「连不上？」时广播这个，电脑收到就弹出网络自检窗口
SELFCHECK_MAGIC = "LANREMOTE_SELFCHECK"
ACK_MAGIC = "LANREMOTE_ACK"
PROTOCOL = 1


def build_reply(name: str, tcp_port: int, version: str) -> bytes:
    payload = {
        "lanremote": REPLY_MAGIC,
        "proto": PROTOCOL,
        "name": name,
        "port": tcp_port,
        "version": version,
    }
    return json.dumps(payload, ensure_ascii=False).encode("utf-8")


def build_selfcheck_ack(name: str, note: str = "") -> bytes:
    """自检应答。

    note 带上电脑这边的诊断结论 —— 手表拿到就直接显示，
    不需要在被控端弹窗口（用户要的是全程无窗口）。
    """
    payload = {"lanremote": ACK_MAGIC, "name": name, "note": note}
    return json.dumps(payload, ensure_ascii=False).encode("utf-8")


def build_selfcheck_request() -> bytes:
    payload = {"lanremote": SELFCHECK_MAGIC, "proto": PROTOCOL}
    return json.dumps(payload, ensure_ascii=False).encode("utf-8")


def parse_selfcheck(data: bytes) -> dict | None:
    try:
        obj = json.loads(data.decode("utf-8", errors="replace"))
    except ValueError:
        return None
    if not isinstance(obj, dict):
        return None
    if obj.get("lanremote") != SELFCHECK_MAGIC:
        return None
    return obj


def parse_probe(data: bytes) -> dict | None:
    try:
        obj = json.loads(data.decode("utf-8", errors="replace"))
    except ValueError:
        return None
    if not isinstance(obj, dict):
        return None
    if obj.get("lanremote") != PROBE_MAGIC:
        return None
    return obj


class DiscoveryResponder:
    """监听 UDP 探测包并回复本机地址。"""

    def __init__(self, store, tcp_port: int, version: str = "",
                 port: int = DISCOVERY_PORT, on_probe=None):
        self.store = store
        self.tcp_port = int(tcp_port)
        self.version = version or _APP_VERSION
        self.port = int(port)
        self.on_probe = on_probe          # 可选回调，用于界面提示"有人在找这台电脑"
        # 收到「请打开网络自检」的请求时调用；由上层决定怎么弹窗口
        self.on_selfcheck = None
        # 最近一次自检的结论，随 ACK 回给手表
        self._selfcheck_note = ""
        self.selfcheck_requests = 0
        self._sock: socket.socket | None = None
        self._thread: threading.Thread | None = None
        self._running = False
        self.probes_answered = 0
        self.last_probe_at = 0.0
        self.last_error = ""

    @property
    def running(self) -> bool:
        return self._running and self._thread is not None and self._thread.is_alive()

    def start(self) -> bool:
        if self.running:
            return True
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            # Windows 上 SO_REUSEADDR 允许第二个进程绑同一 UDP 端口，
            # 两个实例会互相抢探测包。改用独占绑定，第二个实例直接失败。
            if hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
                try:
                    sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
                except OSError:
                    pass
            else:
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
            except OSError:
                pass
            sock.bind(("0.0.0.0", self.port))
            sock.settimeout(1.0)
        except OSError as exc:
            self.last_error = f"无法监听 UDP {self.port}: {exc}"
            return False

        self._sock = sock
        self._running = True
        self._thread = threading.Thread(target=self._loop, name="discovery",
                                        daemon=True)
        self._thread.start()
        return True

    def stop(self) -> None:
        self._running = False
        thread = self._thread
        if thread is not None and thread.is_alive():
            thread.join(timeout=3.0)
        self._thread = None
        if self._sock is not None:
            try:
                self._sock.close()
            except OSError:
                pass
            self._sock = None

    def _loop(self) -> None:
        reply = build_reply(str(self.store.get("server_name", "")),
                            self.tcp_port, self.version)
        while self._running:
            try:
                data, addr = self._sock.recvfrom(2048)
            except socket.timeout:
                continue
            except OSError:
                break

            # 先看是不是「请打开网络自检」的请求
            if parse_selfcheck(data) is not None:
                self.selfcheck_requests += 1
                try:
                    self._sock.sendto(
                        build_selfcheck_ack(
                            str(self.store.get("server_name", "")),
                            self._selfcheck_note),
                        addr)
                except OSError:
                    pass
                if self.on_selfcheck:
                    try:
                        # 回调可以返回一段话，随 ACK 一起回给手表
                        produced = self.on_selfcheck(addr[0])
                        if isinstance(produced, str) and produced:
                            self._selfcheck_note = produced[:1500]
                    except Exception:
                        pass
                continue

            probe = parse_probe(data)
            if probe is None:
                continue

            self.probes_answered += 1
            self.last_probe_at = time.time()
            # 单播回复给发起方
            try:
                self._sock.sendto(reply, addr)
            except OSError as exc:
                self.last_error = str(exc)
                continue

            if self.on_probe:
                try:
                    self.on_probe(addr[0])
                except Exception:
                    pass

    def status(self) -> dict:
        return {
            "running": self.running,
            "port": self.port,
            "answered": self.probes_answered,
            "last_probe": self.last_probe_at,
            "selfcheck_requests": self.selfcheck_requests,
            "error": self.last_error,
        }


def local_broadcasts() -> list[str]:
    """本机各网段可能用到的**定向广播地址**。

    为什么要算这个：很多路由器和安卓版本**不转发 255.255.255.255**
    （受限广播），只转发本网段的定向广播（如 192.168.10.255）。
    只发受限广播时，表现就是"搜不到电脑"。

    纯标准库拿不到 Windows 的掩码，所以：
      * 先取本机所有局域网 IPv4
      * 对每个地址按 /24 推出 x.y.z.255
    家用/办公网绝大多数是 /24，这一招够用；剩下的靠受限广播兜底。
    """
    out: list[str] = []
    try:
        from .server import lan_ips
        ips = lan_ips()
    except Exception:                                     # noqa: BLE001
        ips = []
    for ip in ips:
        if ip.startswith("169.254."):
            continue                        # APIPA/链路本地，发了也没用
        parts = ip.split(".")
        if len(parts) != 4:
            continue
        directed = ".".join(parts[:3] + ["255"])
        if directed not in out:
            out.append(directed)
    return out


# 探测包重发的时间点（秒，相对开始）。
# UDP 广播是尽力而为，只发一次的话丢一个包就整轮搜不到。
_PROBE_AT = (0.0, 0.35, 0.9, 1.8, 3.0)


def discover(timeout: float = 2.5, port: int = DISCOVERY_PORT,
             broadcast: str = "255.255.255.255") -> list[dict]:
    """客户端侧：广播探测并收集回复（手表端用的是同一套逻辑）。

    会**重发多轮**并同时打受限广播和本网段定向广播 —— 原来只发一次、
    而且只发 255.255.255.255，广播一丢或路由器不转发受限广播就搜不到。
    """
    results: list[dict] = []
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        sock.settimeout(0.2)
        probe = json.dumps({"lanremote": PROBE_MAGIC,
                            "proto": PROTOCOL}).encode()

        targets: list[str] = []
        for one in (broadcast,) + tuple(local_broadcasts()):
            if one and one not in targets:
                targets.append(one)

        start = time.time()
        deadline = start + timeout
        next_probe = 0

        while time.time() < deadline:
            # 到点就重发一轮
            elapsed = time.time() - start
            while next_probe < len(_PROBE_AT) and elapsed >= _PROBE_AT[next_probe]:
                for target in targets:
                    try:
                        sock.sendto(probe, (target, port))
                    except OSError:
                        # 某个广播地址不可用时跳过，不影响其他目标
                        pass
                next_probe += 1

            try:
                data, addr = sock.recvfrom(2048)
            except socket.timeout:
                continue
            except OSError:
                break
            try:
                obj = json.loads(data.decode("utf-8", errors="replace"))
            except ValueError:
                continue
            if not isinstance(obj, dict):
                continue
            if obj.get("lanremote") != REPLY_MAGIC:
                continue
            results.append({
                "host": addr[0],
                "port": int(obj.get("port") or 0),
                "name": obj.get("name") or "",
                "version": obj.get("version") or "",
            })
    finally:
        sock.close()
    # 按主机去重
    seen = set()
    unique = []
    for item in results:
        if item["host"] in seen:
            continue
        seen.add(item["host"])
        unique.append(item)
    return unique