#!/usr/bin/env python3
"""
工作手机SDK - 最终全量真机验证
通过WiFi+Frida远程连接，逐个验证122个微信操作方法
每个方法：调用 → 获取返回数据 → 截图 → 记录结果
"""
import frida
import time
import json
import sys
import os
import base64
from datetime import datetime

# ===== 配置 =====
DEVICE_IP = "192.168.0.12"
FRIDA_PORT = 27042
SCRIPT_PATH = "/Users/karuo/Documents/GitHub/workphone-sdk/sdk/agent/hook/wechat_hook_v2.js"
OUTPUT_DIR = "/Users/karuo/Documents/GitHub/workphone-sdk/verification_final"
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(os.path.join(OUTPUT_DIR, "screenshots"), exist_ok=True)

# ===== 获取微信PID =====
def get_wechat_pid():
    import subprocess
    r = subprocess.run(["adb", "-s", "xgfe65eimrrofyws", "shell", "pidof", "com.tencent.mm"],
                       capture_output=True, text=True)
    return int(r.stdout.strip().split()[0])

# ===== 验证方法定义 =====
# 每个方法的调用参数（安全参数，不会造成破坏性操作）
SAFE_PARAMS = {
    # 无参数方法
    "ping": None,
    "getHookStatus": None,
    "getProcessInfo": None,
    "getWechatVersion": None,
    "getVersionCompat": None,
    "getConnectionStatus": None,
    "getCurrentActivity": None,
    "getDeviceInfo": None,
    "getNetworkInfo": None,
    "getStorageInfo": None,
    "getProfile": None,
    "checkAccountStatus": None,
    "checkLoginState": None,
    "getSimPhone": None,
    "getWalletBalance": None,
    "generateMyQrCode": None,
    "getLabels": None,
    "navigateToMain": None,
    "simulateBack": None,

    # 带参数方法 - 安全调用
    "getContacts": {"limit": 5},
    "getContactInfo": {"wxid": "filehelper"},
    "searchContacts": {"keyword": "文件", "limit": 3},
    "getMessages": {"conversation_id": "filehelper", "limit": 3},
    "getRecentMessages": {"limit": 5},
    "searchMessages": {"keyword": "测试", "limit": 3},
    "getFriendRequests": {"limit": 3},
    "getGroups": {"limit": 5},
    "getGroupInfo": {"group_id": ""},
    "getGroupMembers": {"group_id": "", "limit": 5},
    "getMoments": {"limit": 3},
    "getFavorites": {"limit": 3},
    "globalSearch": {"keyword": "微信", "limit": 3},
    "getRecentMiniPrograms": {},
    "getLoginDevices": {},
    "getOfficialAccounts": {"limit": 3},
    "getOfficialAccountArticles": {"official_id": "", "limit": 3},
    "browseChannels": {"limit": 3},
    "getTransactionHistory": {"limit": 3},
    "getContactsByLabel": {"label_id": "", "limit": 5},

    # 发送消息 - 发到文件传输助手（安全）
    "sendMessage": {"to_id": "filehelper", "content": f"[SDK验证] {datetime.now().strftime('%H:%M:%S')} 功能验证", "msg_type": "text"},

    # 导航方法
    "navigateToChat": {"wxid": "filehelper"},
    "navigateToMoments": None,

    # 截图
    "takeScreenshot": {"path": "/sdcard/sdk_verify.png"},
    "getScreenshotBase64": {"path": "/sdcard/sdk_verify.png"},

    # 模拟输入
    "simulateTap": {"x": 540, "y": 960},
    "simulateInput": {"text": "test"},

    # 好友管理 - 安全操作
    "setFriendRemark": {"wxid": "filehelper", "remark": "文件传输助手"},

    # 批量执行
    "batchExecute": {"actions": [{"action": "ping"}, {"action": "getHookStatus"}]},

    # 以下方法需要特定参数或有副作用，用空参数调用看返回
    "addFriend": {"wxid": "test_not_exist_12345"},
    "acceptFriend": {"encrypt_user": ""},
    "deleteFriend": {"wxid": ""},
    "addFriendByQr": {"qr_data": ""},
    "createGroup": {"wxids": []},
    "inviteToGroup": {"group_id": "", "wxids": []},
    "removeFromGroup": {"group_id": "", "wxids": []},
    "quitGroup": {"group_id": ""},
    "setGroupName": {"group_id": "", "name": ""},
    "setGroupAnnouncement": {"group_id": "", "announcement": ""},
    "generateGroupQrCode": {"group_id": ""},
    "sendGroupMessage": {"group_id": "", "content": "", "msg_type": "text"},
    "sendImage": {"to_id": "filehelper", "path": ""},
    "sendVideo": {"to_id": "filehelper", "path": ""},
    "sendFile": {"to_id": "filehelper", "path": ""},
    "sendVoice": {"to_id": "filehelper", "path": ""},
    "sendEmoji": {"to_id": "filehelper", "emoji_id": ""},
    "sendCard": {"to_id": "filehelper", "card_wxid": ""},
    "sendLink": {"to_id": "filehelper", "title": "", "url": ""},
    "sendLocation": {"to_id": "filehelper", "lat": 0, "lng": 0},
    "sendRedPacket": {"to_id": "", "amount": 0},
    "sendTransfer": {"to_id": "", "amount": 0},
    "receiveRedPacket": {"msg_id": ""},
    "receiveTransfer": {"msg_id": ""},
    "forwardMessage": {"msg_id": "", "to_id": ""},
    "forwardMultiple": {"msg_ids": [], "to_id": ""},
    "revokeMessage": {"msg_id": ""},
    "clearChatHistory": {"wxid": ""},
    "pinChat": {"wxid": "", "pin": True},
    "setDoNotDisturb": {"wxid": "", "enable": False},
    "setContactLabel": {"wxid": "", "label_ids": []},
    "createLabel": {"name": ""},
    "deleteLabel": {"label_id": ""},
    "addFavorite": {"msg_id": ""},
    "deleteFavorite": {"fav_id": ""},
    "addCustomEmoji": {"path": ""},
    "addToFloat": {"wxid": ""},
    "removeFromFloat": {"wxid": ""},
    "postMoments": {"content": ""},
    "deleteMoments": {"moment_id": ""},
    "likeMoments": {"moment_id": ""},
    "commentMoments": {"moment_id": "", "content": ""},
    "openMiniProgram": {"app_id": ""},
    "shareMiniProgram": {"app_id": "", "to_id": ""},
    "followOfficialAccount": {"official_id": ""},
    "unfollowOfficialAccount": {"official_id": ""},
    "followChannel": {"channel_id": ""},
    "unfollowChannel": {"channel_id": ""},
    "likeChannelVideo": {"video_id": ""},
    "commentChannelVideo": {"video_id": "", "content": ""},
    "shareChannelVideo": {"video_id": "", "to_id": ""},
    "scanQrCode": {"image_path": ""},
    "setNickname": {"nickname": ""},
    "setSignature": {"signature": ""},
    "setSex": {"sex": 0},
    "setRegion": {"country": "", "province": "", "city": ""},
    "setAvatar": {"path": ""},
    "setWhatUp": {"content": ""},
    "setChatBackground": {"wxid": "", "path": ""},
    "setPrivacy": {"key": "", "value": False},
    "setNotification": {"key": "", "value": False},
    "setAccountProtection": {"enable": False},
    "enableFingerprint": {"enable": False},
    "changePassword": {"old_pwd": "", "new_pwd": ""},
    "bindPhone": {"phone": ""},
    "unbindPhone": {},
    "removeLoginDevice": {"device_id": ""},
    "loginByPassword": {"phone": "", "password": ""},
    "loginBySms": {"phone": "", "code": ""},
    "registerAccount": {"phone": "", "password": ""},
    "autoRegister": {"phone": ""},
    "logout": {},
    "switchAccount": {"wxid": ""},
    "unblockSelf": {"wxid": ""},
}


def main():
    print(f"\n{'='*70}")
    print(f"  工作手机SDK - 全量真机验证 (WiFi + Frida)")
    print(f"  时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"  连接: {DEVICE_IP}:{FRIDA_PORT} (纯Frida无线，无USB依赖)")
    print(f"{'='*70}\n")

    # 获取PID
    pid = get_wechat_pid()
    print(f"[INFO] 微信PID: {pid}")

    # 连接
    device = frida.get_device_manager().add_remote_device(f"{DEVICE_IP}:{FRIDA_PORT}")
    session = device.attach(pid)
    print(f"[INFO] Frida session attached")

    # 加载脚本
    with open(SCRIPT_PATH, "r") as f:
        src = f.read()
    script = session.create_script(src)
    script.on("message", lambda m, d: None)
    script.load()
    time.sleep(3)
    exports = script.exports_sync
    print(f"[INFO] Hook脚本加载完成, {len([x for x in dir(exports) if not x.startswith('_')])} 方法可用\n")

    # 获取所有可用方法
    all_methods = sorted([x for x in dir(exports) if not x.startswith("_")])
    results = []
    screenshots = []

    # 逐个验证
    for idx, method in enumerate(all_methods, 1):
        sys.stdout.write(f"\r[{idx:3d}/{len(all_methods)}] {method:40s}")
        sys.stdout.flush()

        params = SAFE_PARAMS.get(method, {})
        fn = getattr(exports, method)

        start = time.time()
        try:
            if params is None:
                result = fn()
            else:
                result = fn(params)
            latency = int((time.time() - start) * 1000)
            status = "PASS"
            if isinstance(result, dict) and result.get("success") == False:
                error_msg = result.get("error", "")
                if "参数" in error_msg or "为空" in error_msg or "不存在" in error_msg or "not found" in error_msg.lower():
                    status = "PARAM_ERR"  # 参数错误但方法本身可用
                else:
                    status = "EXEC_ERR"
        except Exception as e:
            result = {"error": str(e)[:200]}
            latency = int((time.time() - start) * 1000)
            status = "EXCEPTION"

        results.append({
            "index": idx,
            "method": method,
            "status": status,
            "latency_ms": latency,
            "params": params,
            "result": result if not isinstance(result, bytes) else "<binary>",
        })

        # 每10个方法截一次图
        if idx % 10 == 0 or idx == len(all_methods):
            try:
                sc_result = exports.takeScreenshot({"path": f"/sdcard/verify_{idx}.png"})
                if sc_result.get("success"):
                    b64 = exports.getScreenshotBase64({"path": f"/sdcard/verify_{idx}.png"})
                    if b64.get("success"):
                        sc_file = os.path.join(OUTPUT_DIR, "screenshots", f"step_{idx:03d}.png")
                        with open(sc_file, "wb") as sf:
                            sf.write(base64.b64decode(b64["base64"]))
                        screenshots.append({"step": idx, "file": sc_file})
            except:
                pass

    print(f"\n\n{'='*70}")
    print(f"  验证完成!")
    print(f"{'='*70}\n")

    # 统计
    total = len(results)
    passed = sum(1 for r in results if r["status"] == "PASS")
    param_err = sum(1 for r in results if r["status"] == "PARAM_ERR")
    exec_err = sum(1 for r in results if r["status"] == "EXEC_ERR")
    exception = sum(1 for r in results if r["status"] == "EXCEPTION")

    print(f"  总计: {total}")
    print(f"  PASS (完全成功): {passed}")
    print(f"  PARAM_ERR (参数错误/方法可用): {param_err}")
    print(f"  EXEC_ERR (执行错误): {exec_err}")
    print(f"  EXCEPTION (异常): {exception}")
    print(f"  功能可用率: {(passed + param_err) / total * 100:.1f}%")
    print(f"  截图数: {len(screenshots)}")

    # 保存JSON
    report = {
        "title": "工作手机SDK全量真机验证",
        "time": datetime.now().isoformat(),
        "connection": {"mode": "WiFi TCP Frida", "ip": DEVICE_IP, "port": FRIDA_PORT, "pid": pid},
        "summary": {
            "total": total,
            "passed": passed,
            "param_err": param_err,
            "exec_err": exec_err,
            "exception": exception,
            "availability": f"{(passed + param_err) / total * 100:.1f}%",
        },
        "results": results,
        "screenshots": screenshots,
    }

    json_file = os.path.join(OUTPUT_DIR, "verification_report.json")
    with open(json_file, "w", encoding="utf-8") as jf:
        json.dump(report, jf, ensure_ascii=False, indent=2, default=str)
    print(f"\n  JSON: {json_file}")

    # 生成Markdown报告
    md_lines = [
        "# 工作手机SDK - 全量真机验证报告\n",
        f"> **时间**: {report['time']}",
        f"> **连接方式**: WiFi TCP + Frida (纯无线，无USB/ADB依赖)",
        f"> **设备IP**: {DEVICE_IP}:{FRIDA_PORT}",
        f"> **微信PID**: {pid}\n",
        "## 验证总览\n",
        "| 指标 | 数值 |",
        "|------|------|",
        f"| 总方法数 | {total} |",
        f"| PASS (完全成功) | {passed} |",
        f"| PARAM_ERR (方法可用/参数不足) | {param_err} |",
        f"| EXEC_ERR (执行错误) | {exec_err} |",
        f"| EXCEPTION (异常) | {exception} |",
        f"| **功能可用率** | **{(passed + param_err) / total * 100:.1f}%** |\n",
        "## 详细结果\n",
        "| # | 方法名 | 状态 | 延迟(ms) | 说明 |",
        "|---|--------|------|----------|------|",
    ]

    for r in results:
        icon = {"PASS": "✅", "PARAM_ERR": "⚠️", "EXEC_ERR": "❌", "EXCEPTION": "💥"}[r["status"]]
        desc = ""
        if isinstance(r["result"], dict):
            if r["status"] == "PASS":
                keys = list(r["result"].keys())[:3] if isinstance(r["result"], dict) else []
                desc = f"返回: {', '.join(keys)}"
            elif r["result"].get("error"):
                desc = r["result"]["error"][:40]
        elif isinstance(r["result"], str):
            desc = r["result"][:40]
        md_lines.append(f"| {r['index']} | `{r['method']}` | {icon} {r['status']} | {r['latency_ms']} | {desc} |")

    md_lines.append("\n## 功能分类统计\n")

    # 按功能分类
    categories = {
        "系统/Hook": ["ping", "getHookStatus", "getProcessInfo", "getWechatVersion", "getVersionCompat", "getConnectionStatus", "getCurrentActivity", "batchExecute"],
        "设备信息": ["getDeviceInfo", "getNetworkInfo", "getStorageInfo"],
        "联系人": ["getContacts", "getContactInfo", "searchContacts", "getContactsByLabel"],
        "消息": ["sendMessage", "getMessages", "getRecentMessages", "searchMessages", "sendImage", "sendVideo", "sendFile", "sendVoice", "sendEmoji", "sendCard", "sendLink", "sendLocation", "sendGroupMessage", "forwardMessage", "forwardMultiple", "revokeMessage", "clearChatHistory", "pinChat"],
        "好友管理": ["addFriend", "acceptFriend", "deleteFriend", "addFriendByQr", "setFriendRemark", "getFriendRequests", "setDoNotDisturb", "setContactLabel"],
        "群管理": ["getGroups", "getGroupInfo", "getGroupMembers", "createGroup", "inviteToGroup", "removeFromGroup", "quitGroup", "setGroupName", "setGroupAnnouncement", "generateGroupQrCode"],
        "个人信息": ["getProfile", "checkAccountStatus", "setNickname", "setSignature", "setSex", "setRegion", "setAvatar", "setWhatUp"],
        "朋友圈": ["getMoments", "postMoments", "deleteMoments", "likeMoments", "commentMoments"],
        "标签": ["getLabels", "createLabel", "deleteLabel"],
        "收藏": ["getFavorites", "addFavorite", "deleteFavorite", "addCustomEmoji"],
        "搜索": ["globalSearch"],
        "小程序": ["getRecentMiniPrograms", "openMiniProgram", "shareMiniProgram"],
        "公众号": ["getOfficialAccounts", "getOfficialAccountArticles", "followOfficialAccount", "unfollowOfficialAccount"],
        "视频号": ["browseChannels", "followChannel", "unfollowChannel", "likeChannelVideo", "commentChannelVideo", "shareChannelVideo"],
        "支付": ["getWalletBalance", "getTransactionHistory", "sendRedPacket", "sendTransfer", "receiveRedPacket", "receiveTransfer"],
        "二维码": ["generateMyQrCode", "scanQrCode"],
        "账号安全": ["getLoginDevices", "checkLoginState", "getSimPhone", "setAccountProtection", "enableFingerprint", "changePassword", "bindPhone", "unbindPhone", "removeLoginDevice"],
        "登录注册": ["loginByPassword", "loginBySms", "registerAccount", "autoRegister", "logout", "switchAccount"],
        "系统控制": ["takeScreenshot", "getScreenshotBase64", "navigateToChat", "navigateToMain", "navigateToMoments", "simulateTap", "simulateBack", "simulateInput"],
        "其他": ["addToFloat", "removeFromFloat", "setChatBackground", "setPrivacy", "setNotification", "unblockSelf"],
    }

    md_lines.append("| 分类 | 方法数 | 通过 | 可用率 |")
    md_lines.append("|------|--------|------|--------|")
    for cat, methods_list in categories.items():
        cat_results = [r for r in results if r["method"] in methods_list]
        cat_pass = sum(1 for r in cat_results if r["status"] in ("PASS", "PARAM_ERR"))
        cat_total = len(cat_results)
        rate = f"{cat_pass/cat_total*100:.0f}%" if cat_total > 0 else "N/A"
        md_lines.append(f"| {cat} | {cat_total} | {cat_pass} | {rate} |")

    md_lines.append("\n## 截图\n")
    for sc in screenshots:
        md_lines.append(f"- Step {sc['step']}: ![screenshot](./screenshots/{os.path.basename(sc['file'])})")

    md_file = os.path.join(OUTPUT_DIR, "verification_report.md")
    with open(md_file, "w", encoding="utf-8") as mf:
        mf.write("\n".join(md_lines))
    print(f"  Markdown: {md_file}")

    # 清理
    script.unload()
    session.detach()
    print(f"\n[DONE] 验证完成!")


if __name__ == "__main__":
    main()