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 discover(timeout: float = 2.5, port: int = DISCOVERY_PORT,
             broadcast: str = "255.255.255.255") -> list[dict]:
    """客户端侧：广播探测并收集回复（手表端用的是同一套逻辑）。"""
    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.4)
        probe = json.dumps({"lanremote": PROBE_MAGIC, "proto": PROTOCOL}).encode()
        sock.sendto(probe, (broadcast, port))
        deadline = time.time() + timeout
        while time.time() < deadline:
            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 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