iOSAI/Module/DeviceInfo.py

import json
import os
import socket
import threading
import time
import subprocess
from typing import Dict
import tidevice
import wda
from tidevice import Usbmux, ConnectionType
from Entity.DeviceModel import DeviceModel
from Entity.Variables import WdaAppBundleId, wdaFunctionPort
from Module.FlaskSubprocessManager import FlaskSubprocessManager
from Module.IOSActivator import IOSActivator
from Utils.LogManager import LogManager


class DeviceInfo:
    _instance = None
    _instance_lock = threading.Lock()

    # 离线宽限期（保持你原来的数值）
    REMOVE_GRACE_SEC = 5.0

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            with cls._instance_lock:
                if not cls._instance:
                    cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self) -> None:
        if getattr(self, "_initialized", False):
            return

        self._lock = threading.RLock()
        self._models: Dict[str, DeviceModel] = {}
        self._manager = FlaskSubprocessManager.get_instance()
        self.screenPort = 9110

        # 设备心跳时间
        self._last_seen: Dict[str, float] = {}

        # iproxy 子进程：udid -> Popen
        self._iproxy_process: Dict[str, subprocess.Popen] = {}

        # iproxy HTTP 健康检查失败次数：udid -> 连续失败次数
        self._iproxy_fail_count: Dict[str, int] = {}

        # Windows 下隐藏子进程窗口（给 iproxy 用）
        self._creationflags = 0
        self._startupinfo = None
        if os.name == "nt":
            try:
                # type: ignore[attr-defined]
                self._creationflags = subprocess.CREATE_NO_WINDOW
            except Exception:
                self._creationflags = 0

            si = subprocess.STARTUPINFO()
            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            si.wShowWindow = 0  # SW_HIDE
            self._startupinfo = si

        LogManager.info("DeviceInfo 初始化完成", udid="system")
        print("[Init] DeviceInfo 初始化完成")
        self._initialized = True

    # ==========================
    # 主循环
    # ==========================
    def listen(self):
        LogManager.method_info("进入主循环", "listen", udid="system")
        print("[Listen] 开始监听设备上下线...")

        while True:
            try:
                usb = Usbmux().device_list()
                # 只看 USB 连接的设备
                online = {d.udid for d in usb if d.conn_type == ConnectionType.USB}
            except Exception as e:
                LogManager.warning(f"[device_list] 异常：{e}", udid="system")
                time.sleep(1)
                continue

            now = time.time()

            # 当前已知的设备（本轮循环开始时）
            with self._lock:
                known = set(self._models.keys())
                current_count = len(self._models)

            # 1. 处理在线设备
            for udid in online:
                # 更新心跳时间
                self._last_seen[udid] = now

                # 新设备但数量已达上限
                if udid not in known and current_count >= 6:
                    print(f"[Add] 设备数量已达 6 台，忽略新设备: {udid}")
                    LogManager.info(
                        "[Add] 设备数量已达上限(6)，忽略新设备",
                        udid=udid,
                    )
                    continue

                # 已经在列表里的设备，跳过添加流程
                if udid in known:
                    continue

                # 只对新发现的设备做一次信任检查
                try:
                    if not self._is_trusted(udid):
                        LogManager.info(
                            "[Add] 设备未信任或未就绪，跳过本轮添加",
                            udid=udid,
                        )
                        print(f"[Add] 设备未信任或未就绪，跳过: {udid}")
                        continue
                except Exception as e:
                    LogManager.warning(
                        f"[Add] 检测设备 {udid} 信任状态异常: {e}",
                        udid=udid,
                    )
                    print(f"[Add] 检测设备 {udid} 信任状态异常: {e}")
                    continue

                # 二次确认数量上限
                with self._lock:
                    if len(self._models) >= 6:
                        print(f"[Add] 二次检查: 设备数量已达 6 台，忽略新设备: {udid}")
                        LogManager.info(
                            "[Add] 二次检查数量上限，忽略新设备",
                            udid=udid,
                        )
                        continue

                # 真正添加设备
                try:
                    self._add_device(udid)
                    current_count += 1
                except Exception as e:
                    LogManager.warning(
                        f"[Add] 处理设备 {udid} 异常: {e}",
                        udid=udid,
                    )
                    print(f"[Add] 处理设备 {udid} 异常: {e}")

            # 2. 处理可能离线的设备（只看本轮开始时 known 里的）
            for udid in list(known):
                if udid not in online:
                    last = self._last_seen.get(udid, 0)
                    if now - last > self.REMOVE_GRACE_SEC:
                        try:
                            self._remove_device(udid)
                        except Exception as e:
                            LogManager.method_error(
                                f"移除失败：{e}",
                                "listen",
                                udid=udid,
                            )
                            print(f"[Remove] 移除失败 {udid}: {e}")

            # 3. iproxy 看门狗（进程 + HTTP 探活）
            try:
                self._check_iproxy_health()
            except Exception as e:
                LogManager.warning(
                    f"[iproxy] 看门狗异常: {e}",
                    udid="system",
                )
                print(f"[iproxy] 看门狗异常: {e}")

            time.sleep(1)

    # 判断设备是否信任
    def _is_trusted(self, udid: str) -> bool:
        try:
            d = tidevice.Device(udid)
            _ = d.product_version
            return True
        except Exception as e:
            msg = str(e)
            if "NotTrusted" in msg or "Please trust" in msg or "InvalidHostID" in msg:
                print(f"[Trust] 设备未信任，udid={udid}, err={msg}")
                return False

            print(f"[Trust] 检测信任状态出错，当作未信任处理 udid={udid}, err={msg}")
            return False

    # ==========================
    # 添加设备
    # ==========================
    def _add_device(self, udid: str):
        with self._lock:
            if udid in self._models:
                print(f"[Add] 已存在，跳过: {udid}")
                return
            print(f"[Add] 新增设备 {udid}")

        # 判断 iOS 版本
        try:
            t = tidevice.Device(udid)
            version_major = float(t.product_version.split(".")[0])
        except Exception as e:
            print(f"[Add] 获取系统版本失败 {udid}: {e}")
            version_major = 0

        # 分配投屏端口 & 写入模型
        with self._lock:
            self.screenPort += 1
            screen_port = self.screenPort

            model = DeviceModel(
                deviceId=udid,
                screenPort=screen_port,
                width=0,
                height=0,
                scale=0,
                type=1,
            )
            self._models[udid] = model

        print(f"[Add] 新设备完成 {udid}, screenPort={screen_port}")
        self._manager_send()

        # 启动 iproxy（投屏转发）
        try:
            self._start_iproxy(udid, screen_port)
        except Exception as e:
            print(f"[iproxy] 启动失败 {udid}: {e}")
            LogManager.warning(f"[iproxy] 启动失败: {e}", udid=udid)

        # 启动 WDA
        if version_major >= 17.0:
            threading.Thread(
                target=IOSActivator().activate_ios17,
                args=(udid, self._on_wda_ready),
                daemon=True,
            ).start()
        else:
            try:
                tidevice.Device(udid).app_start(WdaAppBundleId)
            except Exception as e:
                print(f"[Add] 使用 tidevice 启动 WDA 失败 {udid}: {e}")
                LogManager.warning(
                    f"[Add] 使用 tidevice 启动 WDA 失败: {e}",
                    udid=udid,
                )
            else:
                threading.Thread(
                    target=self._fetch_screen_and_notify,
                    args=(udid,),
                    daemon=True,
                ).start()

    # ==========================
    # WDA 启动回调（iOS17+）
    # ==========================
    def _on_wda_ready(self, udid: str):
        print(f"[WDA] 回调触发，准备获取屏幕信息 udid={udid}")
        time.sleep(1)
        threading.Thread(
            target=self._fetch_screen_and_notify,
            args=(udid,),
            daemon=True,
        ).start()

    # ==========================
    # 通过 WDA 获取屏幕信息
    # ==========================
    def _screen_info(self, udid: str):
        try:
            c = wda.USBClient(udid, wdaFunctionPort)
            size = c.window_size()

            w = int(size.width)
            h = int(size.height)
            s = float(c.scale)

            print(f"[Screen] 成功获取屏幕 {w}x{h} scale={s} {udid}")
            return w, h, s
        except Exception as e:
            print(f"[Screen] 获取屏幕失败: {e} udid={udid}")
            return 0, 0, 0.0

    # ==========================
    # 异步获取屏幕尺寸并通知 Flask
    # ==========================
    def _fetch_screen_and_notify(self, udid: str):
        """
        后台线程里多次尝试通过 WDA 获取屏幕尺寸，
        成功后更新 model 并发一次 snapshot。
        """
        max_retry = 15
        interval = 1.0

        time.sleep(2.0)
        for _ in range(max_retry):
            with self._lock:
                if udid not in self._models:
                    print(f"[Screen] 设备已移除，停止获取屏幕信息 udid={udid}")
                    return

            w, h, s = self._screen_info(udid)
            if w > 0 and h > 0:
                with self._lock:
                    m = self._models.get(udid)
                    if not m:
                        print(f"[Screen] 模型已不存在，无法更新 udid={udid}")
                        return
                    m.width = w
                    m.height = h
                    m.scale = s

                print(f"[Screen] 屏幕信息更新完成，准备推送到 Flask udid={udid}")
                try:
                    self._manager_send()
                except Exception as e:
                    print(f"[Screen] 发送屏幕更新到 Flask 失败 udid={udid}, err={e}")
                return

            time.sleep(interval)

        print(f"[Screen] 多次尝试仍未获取到屏幕信息 udid={udid}")

    # ==========================
    # iproxy 管理
    # ==========================
    def _start_iproxy(self, udid: str, local_port: int):
        iproxy_path = self._find_iproxy()

        p = self._iproxy_process.get(udid)
        if p is not None and p.poll() is None:
            print(f"[iproxy] 已存在运行中的进程，跳过 {udid}")
            return

        args = [
            iproxy_path,
            "-u",
            udid,
            str(local_port),  # 本地端口（投屏）
            "9567",           # 手机端口（go-ios screencast）
        ]

        print(f"[iproxy] 启动进程: {args}")

        proc = subprocess.Popen(
            args,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            creationflags=self._creationflags,
            startupinfo=self._startupinfo,
        )

        self._iproxy_process[udid] = proc

    def _stop_iproxy(self, udid: str):
        p = self._iproxy_process.get(udid)
        if not p:
            return
        try:
            p.terminate()
            try:
                p.wait(timeout=2)
            except Exception:
                p.kill()
        except Exception:
            pass
        self._iproxy_process.pop(udid, None)
        print(f"[iproxy] 已停止 {udid}")

    def _is_iproxy_http_healthy(self, local_port: int, timeout: float = 1.0) -> bool:
        """
        通过向本地 iproxy 转发端口发一个最小的 HTTP 请求，
        来判断隧道是否“活着”：
        - 正常：能在超时时间内读到一些 HTTP 头 / 任意字节；
        - 异常：连接失败、超时、完全收不到字节，都认为不健康。
        """
        try:
            with socket.create_connection(("127.0.0.1", local_port), timeout=timeout) as s:
                s.settimeout(timeout)

                req = b"GET / HTTP/1.0\r\nHost: 127.0.0.1\r\n\r\n"
                s.sendall(req)

                data = s.recv(128)
                if not data:
                    return False

                if data.startswith(b"HTTP/") or b"\r\n" in data:
                    return True

                # 即使不是标准 HTTP 头，只要有返回字节，也说明隧道有响应
                return True

        except (socket.timeout, OSError):
            return False
        except Exception:
            return False

    def _check_iproxy_health(self):
        """
        iproxy 看门狗：
        - 先看进程是否存在 / 已退出；
        - 再做一次 HTTP 层探活；
        - 连续多次失败才重启，避免抖动时频繁重启。
        """
        with self._lock:
            items = list(self._models.items())

        for udid, model in items:
            proc = self._iproxy_process.get(udid)

            # 1) 进程不存在或已退出：直接重启
            if proc is None or proc.poll() is not None:
                msg = f"[iproxy] 进程已退出，准备重启 | udid={udid}"
                print(msg)
                LogManager.warning(msg, "iproxy")

                self._iproxy_fail_count[udid] = 0
                try:
                    self._start_iproxy(udid, model.screenPort)
                except Exception as e:
                    msg = f"[iproxy] 重启失败 | udid={udid} | err={e}"
                    print(msg)
                    LogManager.warning(msg, "iproxy")
                continue

            # 2) 进程还在，做一次 HTTP 探活
            is_ok = self._is_iproxy_http_healthy(model.screenPort)

            if is_ok:
                if self._iproxy_fail_count.get(udid):
                    msg = f"[iproxy] HTTP 探活恢复正常 | udid={udid}"
                    print(msg)
                    LogManager.info(msg, "iproxy")
                self._iproxy_fail_count[udid] = 0
                continue

            # 3) HTTP 探活失败：记录一次失败
            fail = self._iproxy_fail_count.get(udid, 0) + 1
            self._iproxy_fail_count[udid] = fail

            msg = f"[iproxy] HTTP 探活失败 {fail} 次 | udid={udid}"
            print(msg)
            LogManager.warning(msg, "iproxy")

            FAIL_THRESHOLD = 3
            if fail >= FAIL_THRESHOLD:
                msg = f"[iproxy] 连续 {fail} 次 HTTP 探活失败，准备重启 | udid={udid}"
                print(msg)
                LogManager.warning(msg, "iproxy")

                self._iproxy_fail_count[udid] = 0
                try:
                    self._stop_iproxy(udid)
                    self._start_iproxy(udid, model.screenPort)
                except Exception as e:
                    msg = f"[iproxy] HTTP 探活重启失败 | udid={udid} | err={e}"
                    print(msg)
                    LogManager.warning(msg, "iproxy")

    # ==========================
    # 移除设备
    # ==========================
    def _remove_device(self, udid: str):
        print(f"[Remove] 移除设备 {udid}")

        self._stop_iproxy(udid)

        with self._lock:
            self._models.pop(udid, None)
            self._last_seen.pop(udid, None)
            self._iproxy_fail_count.pop(udid, None)

        self._manager_send()

    # ==========================
    # 工具方法
    # ==========================
    def _find_iproxy(self) -> str:
        base = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        name = "iproxy.exe" if os.name == "nt" else "iproxy"
        return os.path.join(base, "resources", "iproxy", name)

    # ==========================
    # 同步数据到 Flask
    # ==========================
    def _manager_send(self):
        try:
            self._send_snapshot_to_flask()
        except Exception:
            try:
                self._manager.start()
            except Exception:
                pass
            try:
                self._send_snapshot_to_flask()
            except Exception:
                pass

    def _send_snapshot_to_flask(self):
        with self._lock:
            devices = [m.toDict() for m in self._models.values()]

        payload = json.dumps({"devices": devices}, ensure_ascii=False)
        port = int(os.getenv("FLASK_COMM_PORT", "34566"))

        with socket.create_connection(("127.0.0.1", port), timeout=1.5) as s:
            s.sendall(payload.encode() + b"\n")

        print(f"[SNAPSHOT] 已发送 {len(devices)} 台设备")