workphone-sdk/sdk/agent/hook/wireless_bridge.py

"""
Frida 无线桥接器 — 连接 FridaManager 与 WebSocket 服务端
=======================================================

核心职责：
- 通过 WiFi TCP 连接远程 Frida（替代 USB）
- 将服务端 WebSocket 指令转换为 Frida RPC 调用
- 将 Frida Hook 事件通过 WebSocket 上报服务端
- 管理 Frida 会话生命周期（连接/断开/重连）

数据流：
  服务端 API → WebSocket → Agent → WirelessBridge → FridaManager → Frida RPC → 微信
  微信 → Frida Hook Event → FridaManager → WirelessBridge → WebSocket → 服务端

技术栈：Frida 16.5.6 + asyncio + WebSocket
"""
import os
import json
import time
import logging
import asyncio
import threading
from typing import Optional, Dict, Any, Callable, List
from datetime import datetime

logger = logging.getLogger(__name__)

# 导入同级模块
from .frida_manager import FridaManager, ConnectionMode
from .hook_executor import HookExecutor, ACTION_TO_RPC, MODULE_NAMES
from .event_reporter import EventReporter


class WirelessBridge:
    """
    无线桥接器 — 将 WebSocket 指令转为 Frida RPC 调用
    
    这是 Agent 端的核心组件，负责：
    1. 初始化 Frida 无线连接（remote 模式，通过 WiFi TCP）
    2. 接收 WebSocket 指令并执行对应的 Frida RPC
    3. 将 Hook 事件回传给 WebSocket
    4. 维护连接健康状态
    """

    def __init__(
        self,
        device_ip: str = "127.0.0.1",
        frida_port: int = 27042,
        mode: ConnectionMode = "remote",
        script_path: Optional[str] = None,
        ws_send_fn: Optional[Callable] = None,
    ):
        self.device_ip = device_ip
        self.frida_port = frida_port
        self.mode = mode

        # 默认脚本路径
        if script_path is None:
            script_path = os.path.join(
                os.path.dirname(os.path.abspath(__file__)),
                "wechat_hook_v2.js"
            )

        # 事件上报器
        self.event_reporter = EventReporter()
        if ws_send_fn:
            self.event_reporter.set_send_fn(ws_send_fn)

        # Frida 管理器（使用 remote 模式通过 WiFi 连接）
        self.frida_mgr = FridaManager(
            mode=mode,
            gadget_host=device_ip,
            gadget_port=frida_port,
            script_path=script_path,
            on_event=self.event_reporter.on_hook_event,
            on_detach=self._on_frida_detach,
            auto_reconnect=True,
            reconnect_interval=5.0,
        )

        # Hook 执行器
        self.hook_executor = HookExecutor(self.frida_mgr)

        # 状态
        self._connected = False
        self._stats = {
            "commands_received": 0,
            "commands_success": 0,
            "commands_failed": 0,
            "events_sent": 0,
            "connect_time": None,
            "last_command_time": None,
        }

    # ---- 连接管理 ----

    def connect(self) -> Dict[str, Any]:
        """
        建立 Frida 无线连接
        
        流程：
        1. 通过 WiFi TCP 连接远程 frida-server
        2. attach 到微信进程
        3. 加载 Hook 脚本
        4. 验证 RPC 可用性
        """
        logger.info(f"正在连接 Frida: {self.device_ip}:{self.frida_port} (模式={self.mode})")

        try:
            success = self.frida_mgr.start()
            if not success:
                return {
                    "success": False,
                    "error": "Frida 连接失败",
                    "device_ip": self.device_ip,
                    "frida_port": self.frida_port,
                }

            self._connected = True
            self._stats["connect_time"] = datetime.now().isoformat()

            # 验证 RPC
            ping_result = self.frida_mgr.call_rpc("ping")
            rpc_ok = ping_result.get("success", False)

            logger.info(f"✅ Frida 无线连接成功 | RPC: {'OK' if rpc_ok else 'FAIL'}")

            return {
                "success": True,
                "device_ip": self.device_ip,
                "frida_port": self.frida_port,
                "mode": self.mode,
                "rpc_available": rpc_ok,
                "ping": ping_result.get("data", ""),
                "supported_actions": len(ACTION_TO_RPC),
                "modules": list(MODULE_NAMES.values()),
            }

        except Exception as e:
            self._connected = False
            logger.error(f"Frida 连接异常: {e}")
            return {"success": False, "error": str(e)}

    def disconnect(self):
        """断开 Frida 连接"""
        self.frida_mgr.stop()
        self._connected = False
        logger.info("Frida 无线连接已断开")

    @property
    def connected(self) -> bool:
        return self._connected and self.frida_mgr.connected

    # ---- 指令执行 ----

    def execute_command(self, command: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行 WebSocket 指令
        
        指令格式：
        {
            "type": "command",
            "action": "send_message",
            "params": {"to_id": "xxx", "content": "hello"},
            "command_id": "cmd_123"
        }
        
        返回：
        {
            "success": true/false,
            "command_id": "cmd_123",
            "channel": "hook",
            "data": {...}
        }
        """
        self._stats["commands_received"] += 1
        self._stats["last_command_time"] = datetime.now().isoformat()

        action = command.get("action", "")
        params = command.get("params", {})
        command_id = command.get("command_id", "")

        if not action:
            self._stats["commands_failed"] += 1
            return {
                "success": False,
                "error": "缺少 action",
                "command_id": command_id,
            }

        if not self.connected:
            self._stats["commands_failed"] += 1
            return {
                "success": False,
                "error": "Frida 未连接",
                "command_id": command_id,
            }

        # 通过 HookExecutor 执行
        result = self.hook_executor.execute(action, params)
        result["command_id"] = command_id

        if result.get("success"):
            self._stats["commands_success"] += 1
        else:
            self._stats["commands_failed"] += 1

        return result

    def execute_batch(self, commands: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """批量执行指令"""
        results = []
        for cmd in commands:
            result = self.execute_command(cmd)
            results.append(result)
        return results

    # ---- WebSocket 消息处理 ----

    def handle_ws_message(self, message: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        处理来自 WebSocket 的消息
        
        消息类型：
        - command: 执行 Frida RPC 指令
        - query: 查询状态
        - ping: 心跳
        - batch: 批量执行
        """
        msg_type = message.get("type", "")

        if msg_type == "command":
            return self.execute_command(message)

        elif msg_type == "batch":
            commands = message.get("commands", [])
            return {
                "success": True,
                "type": "batch_result",
                "results": self.execute_batch(commands),
                "total": len(commands),
            }

        elif msg_type == "query":
            query_type = message.get("query", "status")
            if query_type == "status":
                return self.get_status()
            elif query_type == "actions":
                return {
                    "success": True,
                    "actions": self.hook_executor.get_supported_actions(),
                    "total": len(ACTION_TO_RPC),
                }
            elif query_type == "modules":
                return {
                    "success": True,
                    "modules": MODULE_NAMES,
                }
            else:
                return {"success": False, "error": f"未知查询类型: {query_type}"}

        elif msg_type == "ping":
            return {
                "type": "pong",
                "connected": self.connected,
                "timestamp": datetime.now().isoformat(),
            }

        else:
            return {"success": False, "error": f"未知消息类型: {msg_type}"}

    # ---- 事件回调 ----

    def _on_frida_detach(self, reason: str):
        """Frida 会话断开回调"""
        self._connected = False
        logger.warning(f"Frida 会话断开: {reason}")

        # 通知服务端
        if self.event_reporter._send_fn:
            try:
                self.event_reporter._send_fn({
                    "type": "frida_disconnected",
                    "reason": reason,
                    "timestamp": datetime.now().isoformat(),
                })
            except Exception:
                pass

    def set_ws_send_fn(self, fn: Callable):
        """设置 WebSocket 发送函数"""
        self.event_reporter.set_send_fn(fn)

    # ---- 状态查询 ----

    def get_status(self) -> Dict[str, Any]:
        """获取桥接器完整状态"""
        return {
            "success": True,
            "bridge": {
                "connected": self.connected,
                "device_ip": self.device_ip,
                "frida_port": self.frida_port,
                "mode": self.mode,
                "stats": self._stats,
            },
            "frida": self.frida_mgr.get_status(),
            "hook": self.hook_executor.get_status() if self.connected else {
                "available": False,
                "total_actions": len(ACTION_TO_RPC),
            },
            "events": self.event_reporter.get_stats(),
        }

    def get_supported_actions(self) -> Dict[str, Any]:
        """获取支持的所有操作"""
        return {
            "total": len(ACTION_TO_RPC),
            "modules": MODULE_NAMES,
            "actions": {
                module: [
                    action for action, rpc in ACTION_TO_RPC.items()
                    # 简单按模块前缀分组
                ]
                for module in MODULE_NAMES.keys()
            },
            "action_list": sorted(ACTION_TO_RPC.keys()),
        }


# ============================================================
# § 4  异步包装器（用于 FastAPI / asyncio 环境）
# ============================================================

class AsyncWirelessBridge:
    """
    异步包装器 — 将同步的 WirelessBridge 包装为 async 接口
    用于 FastAPI WebSocket 端点
    """

    def __init__(self, bridge: WirelessBridge):
        self._bridge = bridge
        self._loop = None

    async def connect(self) -> Dict[str, Any]:
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._bridge.connect)

    async def disconnect(self):
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._bridge.disconnect)

    async def execute_command(self, command: Dict[str, Any]) -> Dict[str, Any]:
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._bridge.execute_command, command)

    async def handle_ws_message(self, message: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._bridge.handle_ws_message, message)

    @property
    def connected(self) -> bool:
        return self._bridge.connected

    def get_status(self) -> Dict[str, Any]:
        return self._bridge.get_status()