"""
无线 Agent 增强模块 — 为 WorkPhoneAgent 添加 WiFi Frida 支持
=============================================================

功能：
- 自动检测 Root 状态，选择最佳 Frida 连接模式
- 通过 WiFi TCP 连接 Frida（无需 USB）
- 集成 WirelessBridge 处理服务端指令
- 向服务端注册 Frida 连接信息
- 支持 frida-server（Root）和 frida-gadget（免Root）

使用方式：
  在 agent.py 的 __init__ 中调用 init_wireless()
  或在 Termux 中设置环境变量：
    WP_FRIDA_MODE=remote
    WP_FRIDA_HOST=127.0.0.1
    WP_FRIDA_PORT=27042

@author 卡若
@version 1.0.0
"""
import os
import json
import time
import logging
import subprocess
import asyncio
from typing import Optional, Dict, Any

logger = logging.getLogger(__name__)


class WirelessAgentMixin:
    """
    无线 Agent 混入类
    
    将此类混入 WorkPhoneAgent，为其添加 WiFi Frida 支持。
    
    新增属性：
    - wireless_bridge: WirelessBridge 实例
    - frida_mode: 当前 Frida 连接模式
    - frida_host: Frida 服务地址
    - frida_port: Frida 服务端口
    """

    def init_wireless(self):
        """
        初始化无线 Frida 连接
        
        优先级：
        1. 环境变量 WP_FRIDA_MODE / WP_FRIDA_HOST / WP_FRIDA_PORT
        2. 自动检测 Root → remote 模式
        3. 非 Root → gadget 模式
        """
        self.frida_mode = os.environ.get("WP_FRIDA_MODE", "auto").strip().lower()
        self.frida_host = os.environ.get("WP_FRIDA_HOST", "127.0.0.1").strip()
        self.frida_port = int(os.environ.get("WP_FRIDA_PORT", "27042"))

        # 自动检测模式
        if self.frida_mode == "auto":
            if self._check_root():
                self.frida_mode = "remote"
                logger.info("✅ Root 检测通过 → 使用 remote 模式（frida-server）")
            else:
                self.frida_mode = "gadget"
                logger.info("⚠️ 无 Root → 使用 gadget 模式（frida-gadget）")

        # 确保 frida-server 正在运行（Root 模式）
        if self.frida_mode == "remote":
            self._ensure_frida_server_running()

        # 初始化 FridaManager（使用 WiFi 模式）
        self._init_frida_wireless()

        # 初始化 WirelessBridge
        self._init_wireless_bridge()

        logger.info(f"📡 无线 Frida 初始化完成: {self.frida_mode} ({self.frida_host}:{self.frida_port})")

    def _check_root(self) -> bool:
        """检查是否有 Root 权限"""
        try:
            result = subprocess.run(
                ["su", "-c", "id"],
                capture_output=True, text=True, timeout=5,
            )
            return "uid=0" in result.stdout
        except Exception:
            return False

    def _ensure_frida_server_running(self):
        """确保 frida-server 正在运行"""
        try:
            # 检查端口是否已监听
            result = subprocess.run(
                ["su", "-c", f"netstat -tlnp 2>/dev/null | grep :{self.frida_port}"],
                capture_output=True, text=True, timeout=5, shell=True,
            )
            if str(self.frida_port) in result.stdout:
                logger.info(f"frida-server 已在运行 (端口 {self.frida_port})")
                return

            # 尝试启动
            frida_bin = os.path.expanduser("~/.frida-server/frida-server")
            if not os.path.exists(frida_bin):
                # 查找其他可能的路径
                for path in [
                    "/data/local/tmp/frida-server",
                    "/data/local/tmp/app_process64",
                    os.path.expanduser("~/.frida-server/app_process64"),
                ]:
                    if os.path.exists(path):
                        frida_bin = path
                        break

            if os.path.exists(frida_bin):
                subprocess.Popen(
                    ["su", "-c", f"nohup {frida_bin} -l 0.0.0.0:{self.frida_port} &"],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
                time.sleep(2)
                logger.info(f"frida-server 已启动 (端口 {self.frida_port})")
            else:
                logger.warning("frida-server 二进制文件不存在，请先执行部署脚本")

        except Exception as e:
            logger.warning(f"frida-server 启动检查失败: {e}")

    def _init_frida_wireless(self):
        """初始化 Frida Manager（无线模式）"""
        try:
            from hook.frida_manager import FridaManager

            # 如果已有 frida_mgr 且已连接，跳过
            if hasattr(self, 'frida_mgr') and self.frida_mgr and self.frida_mgr.connected:
                logger.info("Frida 已连接，跳过重新初始化")
                return

            self.frida_mgr = FridaManager(
                mode=self.frida_mode,
                gadget_host=self.frida_host,
                gadget_port=self.frida_port,
                on_event=self._on_frida_event if hasattr(self, '_on_frida_event') else None,
                auto_reconnect=True,
                reconnect_interval=5.0,
            )

            started = self.frida_mgr.start()
            if started:
                logger.info(f"🔗 Frida 无线连接成功 ({self.frida_mode})")
            else:
                logger.warning("⚠️ Frida 无线连接失败")
                self.frida_mgr = None

        except ImportError:
            logger.info("Frida 模块未安装")
            self.frida_mgr = None
        except Exception as e:
            logger.warning(f"Frida 无线初始化异常: {e}")
            self.frida_mgr = None

    def _init_wireless_bridge(self):
        """初始化 WirelessBridge"""
        try:
            from hook.wireless_bridge import WirelessBridge

            ws_send_fn = None
            if hasattr(self, 'ws') and self.ws:
                async def _ws_send(data):
                    if self.ws and self.connected:
                        await self.ws.send(json.dumps(data))
                ws_send_fn = _ws_send

            self.wireless_bridge = WirelessBridge(
                device_ip=self.frida_host,
                frida_port=self.frida_port,
                mode=self.frida_mode,
                ws_send_fn=ws_send_fn,
            )

            # 如果 frida_mgr 已连接，直接使用
            if self.frida_mgr and self.frida_mgr.connected:
                self.wireless_bridge.frida_mgr = self.frida_mgr
                self.wireless_bridge._connected = True

            logger.info("WirelessBridge 已初始化")

        except ImportError:
            logger.info("WirelessBridge 模块未安装")
            self.wireless_bridge = None
        except Exception as e:
            logger.warning(f"WirelessBridge 初始化异常: {e}")
            self.wireless_bridge = None

    async def handle_wireless_command(self, data: dict) -> Optional[dict]:
        """
        处理无线 Frida 指令
        
        在 _handle_message 中调用，处理 type=command 的消息。
        """
        if not hasattr(self, 'wireless_bridge') or not self.wireless_bridge:
            return None

        msg_type = data.get("type", "")

        if msg_type == "command":
            action = data.get("action", "")

            # 特殊指令：Frida 连接管理
            if action == "frida_connect":
                params = data.get("params", {})
                self.frida_host = params.get("ip", self.frida_host)
                self.frida_port = params.get("port", self.frida_port)
                self.frida_mode = params.get("mode", self.frida_mode)
                self._init_frida_wireless()
                self._init_wireless_bridge()
                return {
                    "success": self.frida_mgr is not None and self.frida_mgr.connected,
                    "command_id": data.get("command_id"),
                    "wechat_attached": getattr(self.frida_mgr, 'connected', False),
                    "supported_actions": 112,
                }

            elif action == "frida_disconnect":
                if self.frida_mgr:
                    self.frida_mgr.stop()
                if self.wireless_bridge:
                    self.wireless_bridge.disconnect()
                return {
                    "success": True,
                    "command_id": data.get("command_id"),
                }

            # 普通 Frida RPC 指令
            return self.wireless_bridge.execute_command(data)

        elif msg_type in ("query", "ping", "batch"):
            return self.wireless_bridge.handle_ws_message(data)

        return None

    async def register_frida_to_server(self):
        """向服务器注册 Frida 连接信息"""
        if not hasattr(self, 'ws') or not self.ws or not self.connected:
            return

        try:
            device_ip = self._get_device_ip()
            register_data = {
                "type": "event",
                "event": "frida_registered",
                "data": {
                    "device_id": self.device_id,
                    "ip": device_ip,
                    "frida_port": self.frida_port,
                    "mode": self.frida_mode,
                    "frida_connected": self.frida_mgr is not None and self.frida_mgr.connected,
                    "android_version": self._get_prop("ro.build.version.release"),
                    "device_model": self._get_prop("ro.product.model"),
                },
            }
            await self.ws.send(json.dumps(register_data))
            logger.info(f"已向服务器注册 Frida 信息: {device_ip}:{self.frida_port}")
        except Exception as e:
            logger.warning(f"Frida 注册失败: {e}")

    def _get_device_ip(self) -> str:
        """获取设备 WiFi IP"""
        try:
            result = subprocess.run(
                ["ip", "route", "get", "1"],
                capture_output=True, text=True, timeout=5,
            )
            parts = result.stdout.strip().split()
            for i, p in enumerate(parts):
                if p == "src" and i + 1 < len(parts):
                    return parts[i + 1]
        except Exception:
            pass

        try:
            import socket
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except Exception:
            return "127.0.0.1"

    def _get_prop(self, prop: str) -> str:
        """获取 Android 系统属性"""
        try:
            result = subprocess.run(
                ["getprop", prop],
                capture_output=True, text=True, timeout=5,
            )
            return result.stdout.strip()
        except Exception:
            return ""