iOSAI/Module/DeviceInfo.py

# -*- coding: utf-8 -*-
import os
import sys
import time
import json
import wda
import threading
import subprocess
from pathlib import Path
from typing import List, Dict, Optional

from tidevice import Usbmux
from Entity.DeviceModel import DeviceModel
from Entity.Variables import WdaAppBundleId
from Module.FlaskSubprocessManager import FlaskSubprocessManager
from Utils.LogManager import LogManager


class Deviceinfo(object):
    def __init__(self):
        self.deviceIndex = 0
        # 投屏端口（本地映射端口起始值，会递增）
        self.screenProxy = 9110
        # 记录 iproxy Popen 进程：[{ "id": udid, "target": Popen }, ...]
        self.pidList: List[Dict] = []
        # 当前已连接的设备（tidevice 的 Device 对象列表）
        self.deviceArray: List = []

        # 子进程通信（向前端发送设备信息）
        self.manager = FlaskSubprocessManager.get_instance()
        # 已发给前端的设备模型列表（用于拔出时发 type=2）
        self.deviceModelList: List[DeviceModel] = []
        # 最大可连接设备限制
        self.maxDeviceCount = 6

        # ===== iproxy：一次性完成 路径定位 + 环境变量配置 + 启动器准备 =====
        try:
            self.iproxy_path = self._iproxy_path()              # 绝对路径
            self.iproxy_dir = self.iproxy_path.parent

            # 1) 配置环境（PATH/DLL），放到初始化里一次性处理
            os.environ["PATH"] = str(self.iproxy_dir) + os.pathsep + os.environ.get("PATH", "")
            try:
                # 仅 Windows 有效；其他平台忽略
                os.add_dll_directory(str(self.iproxy_dir))
            except Exception:
                pass

            # 2) 预构建通用 Popen 参数（隐藏窗口、工作目录、文本模式等）
            self._creationflags = 0x08000000 if os.name == "nt" else 0
            self._popen_kwargs = dict(
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                cwd=str(self.iproxy_dir),
                shell=False,
                text=True,
                encoding="utf-8",
                bufsize=1,
                creationflags=self._creationflags,
            )

            # 3) 准备一个“启动器”（闭包）：仅接受 (udid, local_port, remote_port) 参数
            def _spawn_iproxy(udid: str, local_port: int, remote_port: int = 9100) -> subprocess.Popen:
                args = [str(self.iproxy_path), "-u", udid, str(local_port), str(remote_port)]
                p = subprocess.Popen(args, **self._popen_kwargs)

                # 异步日志转发（可选）
                def _pipe_to_log(name: str, stream):
                    try:
                        for line in iter(stream.readline, ''):
                            s = line.strip()
                            if s:
                                LogManager.info(f"[iproxy {name}] {s}", udid)
                    except Exception:
                        pass

                try:
                    import threading
                    threading.Thread(target=_pipe_to_log, args=("STDOUT", p.stdout), daemon=True).start()
                    threading.Thread(target=_pipe_to_log, args=("STDERR", p.stderr), daemon=True).start()
                except Exception:
                    pass

                return p

            self._spawn_iproxy = _spawn_iproxy  # 保存启动器
            LogManager.info(f"iproxy 启动器已就绪，目录: {self.iproxy_dir}")

        except Exception as e:
            # 没找到 iproxy 也允许实例化成功，但后续启动会失败并给出明确日志
            self.iproxy_path = None
            self.iproxy_dir = None
            self._spawn_iproxy = None
            LogManager.error(f"初始化 iproxy 失败：{e}")

    # ----------------------------
    # 监听设备连接（死循环，内部捕获异常）
    # ----------------------------
    def startDeviceListener(self):

        while True:
            try:
                lists = Usbmux().device_list()
            except Exception as e:
                # 另一台电脑常见：usbmuxd 连接失败（未安装 iTunes/Apple Mobile Device Support）
                LogManager.warning(f"usbmuxd 连接失败: {e}。请确认已安装 iTunes/Apple Mobile Device Support，并在手机上“信任此电脑”")
                time.sleep(2)
                continue

            # 新接入设备
            for device in lists:
                if (device not in self.deviceArray) and (len(self.deviceArray) < self.maxDeviceCount):
                    self.screenProxy += 1
                    try:
                        self.connectDevice(device.udid)
                        self.deviceArray.append(device)
                    except Exception as e:
                        LogManager.error(f"连接设备失败 {device.udid}: {e}", device.udid)

            # 拔出设备处理
            self._removeDisconnected(lists)
            time.sleep(1)

    # ----------------------------
    # 连接单台设备：启动 WDA、读取屏参、通知前端、映射投屏端口
    # ----------------------------
    def connectDevice(self, identifier: str):
        # 1) 连接 WDA（USBClient -> 设备 8100）
        try:
            d = wda.USBClient(identifier, 8100)
            LogManager.info("启动 WDA 成功", identifier)
        except Exception as e:
            LogManager.error(f"启动 WDA 失败，请检查手机是否已信任、WDA 是否正常。错误: {e}", identifier)
            return  # 不抛出到外层，保持监听循环健壮

        # 2) 读取屏幕信息（失败不影响主流程）
        width, height, scale = 0, 0, 1.0
        try:
            size = d.window_size()
            width, height = size.width, size.height
            scale = d.scale
        except Exception as e:
            LogManager.warning(f"读取屏幕信息失败：{e}", identifier)

        # 3) 组装模型并发送给前端
        model = DeviceModel(identifier, self.screenProxy, width, height, scale, type=1)
        self.deviceModelList.append(model)
        try:
            self.manager.send(model.toDict())
        except Exception as e:
            LogManager.warning(f"向前端发送设备模型失败：{e}", identifier)

        # 4) 可选：启动你的 app 并回到桌面
        try:
            d.app_start(WdaAppBundleId)
            d.home()
        except Exception as e:
            LogManager.warning(f"启动/切回桌面失败：{e}", identifier)

        time.sleep(2)

        # 5) 本地端口 -> 设备端口 的映射（投屏：本地 self.screenProxy -> 设备 9100）
        target = self.relayDeviceScreenPort(identifier)
        self.pidList.append({"target": target, "id": identifier})

    # ----------------------------
    # 处理拔出设备：发通知、关掉 iproxy、移出状态
    # ----------------------------
    def _removeDisconnected(self, current_list):
        set1 = set(self.deviceArray)
        set2 = set(current_list)
        difference = list(set1 - set2)  # 在旧集合中但不在新集合中 -> 已拔出

        for i in difference:
            udid = i.udid
            # 1) 通知前端：type = 2
            for a in list(self.deviceModelList):
                if udid == a.deviceId:
                    a.type = 2
                    try:
                        self.manager.send(a.toDict())
                    except Exception as e:
                        LogManager.warning(f"发送下线事件失败：{e}", udid)
                    self.deviceModelList.remove(a)

            # 2) 关掉对应的 iproxy
            for k in list(self.pidList):
                if udid == k["id"]:
                    target = k.get("target")
                    try:
                        if target and target.poll() is None:
                            target.kill()
                    except Exception:
                        pass
                    self.pidList.remove(k)

            # 3) 从已连接集合中移除
            try:
                self.deviceArray.remove(i)
            except Exception:
                pass

    # ----------------------------
    # 根目录与 iproxy 可执行文件定位
    # ----------------------------
    def _base_dir(self) -> Path:
        if getattr(sys, "frozen", False):
            return Path(sys.executable).resolve().parent
        return Path(__file__).resolve().parents[1]  # iOSAI/ 作为根

    def _iproxy_path(self) -> Path:
        exe = "iproxy.exe" if os.name == "nt" else "iproxy"
        base = self._base_dir()
        candidates = [
            base / "resources" / "iproxy" / exe,   # 推荐放置
        ]
        for p in candidates:
            if p.exists():
                return p
        raise FileNotFoundError(f"iproxy not found, tried: {[str(c) for c in candidates]}")

    # ----------------------------
    # 端口映射：仅做“转发端口”这件事（调用已准备好的启动器）
    # ----------------------------
    def relayDeviceScreenPort(self, udid: str) -> Optional[subprocess.Popen]:
        if not self._spawn_iproxy:
            LogManager.error("iproxy 启动器未就绪，无法建立端口映射（初始化时未找到 iproxy）。", udid)
            return None

        try:
            p = self._spawn_iproxy(udid, self.screenProxy, 9100)
            LogManager.info(f"启动 iproxy 成功，本地 {self.screenProxy} -> 设备 9100", udid)
            return p
        except Exception as e:
            LogManager.error(f"启动 iproxy 失败：{e}", udid)
            return None