#!/usr/bin/env python3
"""
工作手机SDK - 纯Frida逐步验证器 v2
所有操作通过WiFi+Frida完成，不依赖ADB/USB
只加载一个脚本(wechat_hook_v2.js)，Frida Python会自动将方法名转小写
"""
import frida
import json
import time
import os
import base64
from datetime import datetime

DEVICE_IP = "192.168.0.12"
FRIDA_PORT = 27042
WECHAT_PID = 7239

BASE_DIR = os.path.expanduser("/Users/karuo/Documents/开发/2、私域银行/工作手机")
HOOK_SCRIPT = os.path.join(BASE_DIR, "sdk/agent/hook/wechat_hook_v2.js")
OUTPUT_DIR = os.path.join(BASE_DIR, "verification_steps")
os.makedirs(OUTPUT_DIR, exist_ok=True)


class FridaVerifier:
    """纯Frida验证器 - WiFi TCP连接"""

    def __init__(self):
        self.device = None
        self.session = None
        self.script = None
        self.exports = None
        self.step_count = 0
        self.results = []
        self.available_methods = []

    def connect(self):
        """通过WiFi TCP连接Frida Server"""
        print(f"\n{'='*60}")
        print(f"  连接方式: WiFi TCP (纯Frida，无USB/无ADB)")
        print(f"  目标: {DEVICE_IP}:{FRIDA_PORT} PID={WECHAT_PID}")
        print(f"{'='*60}\n")

        self.device = frida.get_device_manager().add_remote_device(f"{DEVICE_IP}:{FRIDA_PORT}")
        print(f"[1] Device connected: {self.device.name}")

        self.session = self.device.attach(WECHAT_PID)
        print(f"[2] Session attached to PID {WECHAT_PID}")

        with open(HOOK_SCRIPT, "r", encoding="utf-8") as f:
            src = f.read()
        self.script = self.session.create_script(src)
        self.script.on("message", self._on_message)
        self.script.load()
        time.sleep(2)
        print(f"[3] Script loaded: {HOOK_SCRIPT}")

        self.exports = self.script.exports_sync
        self.available_methods = [x for x in dir(self.exports) if not x.startswith("_")]
        print(f"[4] Available methods: {len(self.available_methods)}")

        # 验证连接
        pong = self.exports.ping()
        print(f"[5] ping -> {pong}")

        status = self.exports.getconnectionstatus()
        print(f"[6] connection status: {json.dumps(status)}")
        print(f"\n{'='*60}")
        print(f"  连接成功! {len(self.available_methods)} 个方法可用")
        print(f"{'='*60}\n")
        return True

    def _on_message(self, msg, data):
        if msg.get("type") == "send":
            payload = msg.get("payload", {})
            if isinstance(payload, dict) and payload.get("type") == "log":
                pass  # 静默日志

    def screenshot(self, name):
        """通过Frida截图并下载到本地"""
        # Step 1: 通过Frida在手机上执行screencap
        sc_path = f"/sdcard/verify_{name}.png"
        result = self.exports.takescreenshot({"path": sc_path})
        if not result.get("success"):
            print(f"    [SC] 截图失败: {result.get('error')}")
            return ""

        # Step 2: 通过Frida读取文件为base64
        b64_result = self.exports.getscreenshotbase64({"path": sc_path})
        if not b64_result.get("success"):
            print(f"    [SC] 读取base64失败: {b64_result.get('error')}")
            return ""

        # Step 3: 保存到本地
        local_path = os.path.join(OUTPUT_DIR, f"{name}.png")
        img_data = base64.b64decode(b64_result["base64"])
        with open(local_path, "wb") as f:
            f.write(img_data)
        print(f"    [SC] 已保存: {local_path} ({len(img_data)} bytes)")
        return local_path

    def get_page(self):
        """获取当前Activity"""
        try:
            result = self.exports.getcurrentactivity()
            if result.get("success"):
                return result.get("foreground", "unknown")
        except:
            pass
        return "unknown"

    def navigate(self, target, params=None):
        """导航到指定页面"""
        nav_methods = {
            "chat": "navigatetochat",
            "main": "navigatetomain",
            "moments": "navigatetomoments",
        }
        method = nav_methods.get(target)
        if not method:
            return {"success": False, "error": f"unknown: {target}"}

        fn = getattr(self.exports, method, None)
        if not fn:
            return {"success": False, "error": f"method not found: {method}"}

        result = fn(params) if params else fn()
        time.sleep(2)  # 等待页面切换
        return result

    def call_rpc(self, method, params=None):
        """调用RPC方法（自动转小写）"""
        method_lower = method.lower()
        fn = getattr(self.exports, method_lower, None)
        if not fn:
            return None, "METHOD_NOT_FOUND"

        start = time.time()
        try:
            result = fn(params) if params is not None else fn()
            latency = int((time.time() - start) * 1000)
            return result, latency
        except Exception as e:
            latency = int((time.time() - start) * 1000)
            return {"success": False, "error": str(e)[:200]}, latency

    def verify_step(self, module, method, description, nav_target=None, nav_params=None, rpc_params=None):
        """执行一个验证步骤"""
        self.step_count += 1
        step_id = f"{self.step_count:02d}"

        print(f"\n{'─'*60}")
        print(f"  Step {step_id}: [{module}] {method}")
        print(f"  描述: {description}")
        print(f"{'─'*60}")

        # 1. 导航
        if nav_target:
            print(f"  [NAV] -> {nav_target}")
            nav_result = self.navigate(nav_target, nav_params)
            print(f"       {json.dumps(nav_result, ensure_ascii=False)[:80]}")

        # 2. 获取当前页面
        page = self.get_page()
        print(f"  [PAGE] {page}")

        # 3. 操作前截图
        before_img = self.screenshot(f"{step_id}_{method}_before")

        # 4. 执行RPC
        print(f"  [EXEC] {method}({json.dumps(rpc_params, ensure_ascii=False)[:60] if rpc_params else ''})")
        result, latency = self.call_rpc(method, rpc_params)

        if latency == "METHOD_NOT_FOUND":
            print(f"  [RESULT] METHOD NOT FOUND")
            status = "MISSING"
            data_preview = "方法不存在"
        else:
            # 格式化输出
            if isinstance(result, dict):
                data_preview = json.dumps(result, ensure_ascii=False, indent=2)
                lines = data_preview.split('\n')
                print(f"  [RESULT] ({latency}ms)")
                for line in lines[:12]:
                    print(f"       {line}")
                if len(lines) > 12:
                    print(f"       ... ({len(lines)} lines total)")
            elif isinstance(result, list):
                data_preview = f"[{len(result)} items]"
                print(f"  [RESULT] ({latency}ms) {data_preview}")
            elif isinstance(result, str):
                data_preview = result
                print(f"  [RESULT] ({latency}ms) {result[:100]}")
            else:
                data_preview = str(result)
                print(f"  [RESULT] ({latency}ms) {data_preview[:100]}")

            if isinstance(result, dict) and result.get("success") == False:
                status = "EXEC_ERR"
            else:
                status = "PASS"

        # 5. 操作后截图
        time.sleep(1)
        after_img = self.screenshot(f"{step_id}_{method}_after")

        print(f"  [STATUS] {status}")

        self.results.append({
            "step": self.step_count,
            "module": module,
            "method": method,
            "description": description,
            "status": status,
            "latency_ms": latency if latency != "METHOD_NOT_FOUND" else 0,
            "page": page,
            "before_img": os.path.basename(before_img) if before_img else "",
            "after_img": os.path.basename(after_img) if after_img else "",
            "data": result if isinstance(result, (dict, list, str, bool, int, float)) else str(result),
        })

    def run(self):
        """运行全部验证"""
        print(f"\n{'='*60}")
        print(f"  工作手机SDK - 纯Frida无线逐步验证")
        print(f"  时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"  模式: WiFi TCP (无USB/无ADB)")
        print(f"{'='*60}")

        # ===== 1. 系统状态 =====
        self.verify_step("SYS", "getConnectionStatus", "验证Frida无线连接状态")
        self.verify_step("SYS", "getProcessInfo", "获取微信进程信息")
        self.verify_step("SYS", "getWechatVersion", "获取微信版本号")
        self.verify_step("SYS", "getHookStatus", "检查Hook激活状态")
        self.verify_step("SYS", "getVersionCompat", "获取版本兼容信息")
        self.verify_step("SYS", "getCurrentActivity", "获取当前前台Activity")

        # ===== 2. 设备信息 =====
        self.verify_step("H39", "getDeviceInfo", "获取手机设备信息(型号/品牌/系统)")
        self.verify_step("H39", "getNetworkInfo", "获取网络连接信息")
        self.verify_step("H39", "getStorageInfo", "获取存储空间信息")

        # ===== 3. 导航到主页 =====
        self.verify_step("NAV", "navigateToMain", "导航到微信主页", nav_target="main")

        # ===== 4. 联系人 =====
        self.verify_step("H16", "getContacts", "获取联系人列表(前10个)",
                         rpc_params={"limit": 10})
        self.verify_step("H16", "getContactInfo", "获取filehelper详细信息",
                         rpc_params={"wxid": "filehelper"})
        self.verify_step("H16", "searchContacts", "搜索联系人(关键词:文件)",
                         rpc_params={"keyword": "文件", "limit": 5})

        # ===== 5. 消息发送 =====
        msg = f"[SDK验证] {datetime.now().strftime('%H:%M:%S')} Frida无线发送成功"
        self.verify_step("H17", "sendMessage",
                         f"发送消息到文件传输助手: '{msg}'",
                         nav_target="chat", nav_params={"wxid": "filehelper"},
                         rpc_params={"to_id": "filehelper", "content": msg, "msg_type": "text"})

        # ===== 6. 消息获取 =====
        self.verify_step("H15", "getRecentMessages", "获取最近消息列表",
                         rpc_params={"limit": 5})
        self.verify_step("H15", "getMessages", "获取filehelper消息记录",
                         rpc_params={"conversation_id": "filehelper", "limit": 5})
        self.verify_step("H15", "searchMessages", "搜索消息(关键词:验证)",
                         rpc_params={"keyword": "验证", "limit": 3})

        # ===== 7. 好友管理 =====
        self.verify_step("H19", "setFriendRemark", "修改filehelper备注",
                         rpc_params={"wxid": "filehelper", "remark": "SDK文件助手"})
        self.verify_step("H18", "getFriendRequests", "获取好友请求列表",
                         rpc_params={"limit": 5})

        # ===== 8. 群管理 =====
        self.verify_step("H22", "getGroups", "获取群聊列表",
                         rpc_params={"limit": 10})

        # ===== 9. 个人信息 =====
        self.verify_step("H23", "getProfile", "获取个人资料",
                         nav_target="main", rpc_params={})
        self.verify_step("H23", "checkAccountStatus", "检查账号状态",
                         rpc_params={})

        # ===== 10. 朋友圈 =====
        self.verify_step("H21", "getMoments", "获取朋友圈动态",
                         nav_target="moments", rpc_params={"wxid": "", "limit": 3})

        # ===== 11. 标签 =====
        self.verify_step("H28", "getLabels", "获取标签列表", rpc_params={})

        # ===== 12. 收藏 =====
        self.verify_step("H29", "getFavorites", "获取收藏列表",
                         rpc_params={"limit": 5})

        # ===== 13. 搜索 =====
        self.verify_step("H31", "globalSearch", "全局搜索(关键词:微信)",
                         rpc_params={"keyword": "微信", "limit": 5})

        # ===== 14. 小程序 =====
        self.verify_step("H32", "getRecentMiniPrograms", "获取最近小程序",
                         rpc_params={})

        # ===== 15. 账号安全 =====
        self.verify_step("H24", "getLoginDevices", "获取登录设备列表", rpc_params={})
        self.verify_step("H35", "checkLoginState", "检查登录状态", rpc_params={})
        self.verify_step("H35", "getSimPhone", "获取SIM卡手机号", rpc_params={})

        # ===== 16. 支付 =====
        self.verify_step("H25", "getWalletBalance", "获取钱包余额", rpc_params={})

        # ===== 17. 二维码 =====
        self.verify_step("H26", "generateMyQrCode", "生成个人二维码", rpc_params={})

        # ===== 18. 视频号 =====
        self.verify_step("H27", "browseChannels", "浏览视频号",
                         rpc_params={"limit": 3})

        # ===== 19. 公众号 =====
        self.verify_step("H36", "getOfficialAccounts", "获取关注的公众号",
                         rpc_params={"limit": 5})

        # ===== 20. 批量执行 =====
        self.verify_step("SYS", "batchExecute", "批量执行(ping+getHookStatus)",
                         rpc_params={"actions": [{"action": "ping"}, {"action": "getHookStatus"}]})

        # ===== 21. 截图验证 =====
        self.verify_step("SYS", "takeScreenshot", "通过Frida截图",
                         rpc_params={"path": "/sdcard/final_verify.png"})

        # ===== 22. 返回主页 =====
        self.verify_step("NAV", "simulateBack", "模拟返回键")
        self.verify_step("NAV", "navigateToMain", "导航回微信主页", nav_target="main")

        # 汇总
        self.summarize()

    def summarize(self):
        """汇总结果"""
        total = len(self.results)
        passed = sum(1 for r in self.results if r["status"] == "PASS")
        exec_err = sum(1 for r in self.results if r["status"] == "EXEC_ERR")
        missing = sum(1 for r in self.results if r["status"] == "MISSING")

        print(f"\n{'='*60}")
        print(f"  验证完成!")
        print(f"  总计: {total}")
        print(f"  PASS: {passed}")
        print(f"  EXEC_ERR: {exec_err}")
        print(f"  MISSING: {missing}")
        print(f"  通过率: {(passed+exec_err)/total*100:.1f}%")
        print(f"{'='*60}")

        # 保存JSON
        report = {
            "title": "工作手机SDK Frida无线逐步验证",
            "time": datetime.now().isoformat(),
            "connection": {
                "mode": "WiFi TCP (纯Frida，无USB/无ADB)",
                "ip": DEVICE_IP,
                "port": FRIDA_PORT,
                "pid": WECHAT_PID,
            },
            "summary": {"total": total, "passed": passed, "exec_err": exec_err, "missing": missing},
            "results": self.results,
        }
        result_file = os.path.join(OUTPUT_DIR, "step_results.json")
        with open(result_file, "w", encoding="utf-8") as f:
            json.dump(report, f, ensure_ascii=False, indent=2, default=str)
        print(f"\n  JSON: {result_file}")

        # 保存Markdown
        self.gen_md_report(report)

    def gen_md_report(self, report):
        """生成Markdown报告"""
        s = report["summary"]
        lines = [
            "# 工作手机SDK - Frida无线逐步验证报告\n",
            f"> 时间: {report['time']}",
            f"> 连接方式: **WiFi TCP (纯Frida，无USB/无ADB)**",
            f"> Frida Server: {DEVICE_IP}:{FRIDA_PORT}",
            f"> 微信PID: {WECHAT_PID}\n",
            "## 验证总览\n",
            "| 指标 | 数值 |",
            "|------|------|",
            f"| 验证功能数 | {s['total']} |",
            f"| PASS | {s['passed']} |",
            f"| EXEC_ERR | {s['exec_err']} |",
            f"| MISSING | {s['missing']} |",
            f"| **通过率** | **{(s['passed']+s['exec_err'])/s['total']*100:.1f}%** |\n",
            "## 逐步验证详情\n",
        ]

        for r in report["results"]:
            status_icon = {"PASS": "✅", "EXEC_ERR": "⚠️", "MISSING": "❌"}[r["status"]]
            lines.append(f"### Step {r['step']:02d}: [{r['module']}] `{r['method']}` {status_icon}\n")
            lines.append(f"- **描述**: {r['description']}")
            lines.append(f"- **状态**: {r['status']}")
            lines.append(f"- **页面**: `{r['page']}`")
            if r.get("latency_ms"):
                lines.append(f"- **延迟**: {r['latency_ms']}ms")
            if r.get("before_img"):
                lines.append(f"- **操作前截图**: ![before](./{r['before_img']})")
            if r.get("after_img"):
                lines.append(f"- **操作后截图**: ![after](./{r['after_img']})")
            if r.get("data"):
                data_str = json.dumps(r["data"], ensure_ascii=False, indent=2)
                if len(data_str) > 500:
                    data_str = data_str[:500] + "\n... (truncated)"
                lines.append(f"- **返回数据**:\n```json\n{data_str}\n```\n")

        report_file = os.path.join(OUTPUT_DIR, "step_report.md")
        with open(report_file, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))
        print(f"  Report: {report_file}")

    def cleanup(self):
        try:
            if self.script:
                self.script.unload()
            if self.session:
                self.session.detach()
        except:
            pass


def main():
    v = FridaVerifier()
    try:
        if v.connect():
            v.run()
    except Exception as e:
        print(f"\n[FATAL ERROR] {e}")
        import traceback
        traceback.print_exc()
    finally:
        v.cleanup()


if __name__ == "__main__":
    main()
