#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
工作手机SDK - 微信全功能Frida验证脚本 v3
基于WCDB直接操作，无感控制微信所有功能
"""
import frida, time, json, subprocess, os
from datetime import datetime

PHONE_IP = "192.168.110.80"
ADB_SERIAL = "192.168.110.80:5555"
FRIDA_PORT = 27042
WECHAT_PID = 9462
BASE = "/Users/karuo/Documents/开发/2、私域银行/工作手机"
HOOK_JS = BASE + "/sdk/agent/hook/wechat_full_control.js"
SC_DIR = BASE + "/开发文档/6、测试/live_verify_20260518/screenshots"
RESULT_FILE = BASE + "/verification_full_v3_20260518.json"

os.makedirs(SC_DIR, exist_ok=True)

def adb_sc(name, wait=0.5):
    time.sleep(wait)
    remote = f"/sdcard/{name}.png"
    local = f"{SC_DIR}/{name}.png"
    subprocess.run(["adb", "-s", ADB_SERIAL, "shell", "screencap", "-p", remote], capture_output=True)
    r = subprocess.run(["adb", "-s", ADB_SERIAL, "pull", remote, local], capture_output=True)
    size = os.path.getsize(local) if os.path.exists(local) else 0
    status = "OK" if size > 10000 else "FAIL(small)"
    print(f"    [截图] {name}.png {size//1024}KB [{status}]")
    return local if size > 10000 else None

print(f"\n{'='*60}")
print(f"工作手机SDK - 微信全功能验证 v3")
print(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
print(f"{'='*60}")

# 连接Frida
print(f"\n[连接] Frida {PHONE_IP}:{FRIDA_PORT} PID={WECHAT_PID}")
dm = frida.get_device_manager()
device = dm.add_remote_device(f"{PHONE_IP}:{FRIDA_PORT}")
session = device.attach(WECHAT_PID)

events = []
def on_msg(m, d):
    if m.get("type") == "send":
        p = m["payload"]
        events.append(p)
        t = p.get("type", "")
        if t in ("ready", "hook_ok", "hook_fail", "real_send_captured", "message_sent", "db_query_ok", "message_inserted"):
            print(f"  [EVENT:{t}] {json.dumps(p, ensure_ascii=False)[:100]}")

with open(HOOK_JS, encoding="utf-8") as f:
    src = f.read()
script = session.create_script(src)
script.on("message", on_msg)
script.load()
print("[等待3秒让Hook初始化...]")
time.sleep(3)
rpc = script.exports_sync

results = {}
passed = 0
failed = 0
ts = datetime.now().strftime("%H%M%S")

def test(name, func, *args):
    global passed, failed
    try:
        r = func(*args) if args else func()
        ok = r.get("success", False) if isinstance(r, dict) else bool(r)
        status = "✅ PASS" if ok else "❌ FAIL"
        print(f"  {status} {name}")
        if not ok and isinstance(r, dict) and r.get("error"):
            print(f"         错误: {r['error'][:80]}")
        if ok: passed += 1
        else: failed += 1
        return r
    except Exception as e:
        print(f"  ❌ FAIL {name} [异常: {str(e)[:60]}]")
        failed += 1
        return {"success": False, "error": str(e)}

# ============================================================
# 模块1：系统基础
# ============================================================
print(f"\n{'='*50}")
print("模块1：系统基础")
print(f"{'='*50}")

r_sys = test("系统信息", rpc.get_system_info)
print(f"    DB路径: {r_sys.get('db', {}).get('path', 'N/A')}")
print(f"    ClassLoader: {r_sys.get('loader', 'N/A')}")

r_db = test("DB实例枚举", rpc.get_db_info)
print(f"    DB实例数: {r_db.get('count', 0)}")
for inst in r_db.get("instances", []):
    print(f"    - {inst.get('path', 'N/A')} [open={inst.get('open')}]")

results["system"] = {"sys_info": r_sys, "db_info": r_db}

# ============================================================
# 模块2：账号信息
# ============================================================
print(f"\n{'='*50}")
print("模块2：账号信息")
print(f"{'='*50}")

r_acc = test("账号信息", rpc.get_account_info)
if r_acc.get("success"):
    acc = r_acc.get("account", {})
    print(f"    昵称: {acc.get('nickname', 'N/A')}")
    print(f"    微信号: {acc.get('wxid', 'N/A')}")
    print(f"    手机: {acc.get('mobile', 'N/A')}")
    print(f"    别名: {acc.get('alias', 'N/A')}")
results["account"] = r_acc

# ============================================================
# 模块3：联系人（完整数据）
# ============================================================
print(f"\n{'='*50}")
print("模块3：联系人（完整数据）")
print(f"{'='*50}")

r_cnt = test("联系人总数", rpc.get_contacts_count)
if r_cnt.get("success"):
    print(f"    真实好友(type=3): {r_cnt.get('type3_friends')} 人")
    print(f"    总联系人记录: {r_cnt.get('total_rcontact')} 条")
    print(f"    群组数量: {r_cnt.get('groups')} 个")

r_contacts = test("联系人列表(前20)", lambda: rpc.get_contacts({"limit": 20, "offset": 0}))
if r_contacts.get("success"):
    print(f"    返回: {r_contacts.get('count', 0)} 条")
    for c in r_contacts.get("contacts", [])[:5]:
        name = c.get("conRemark") or c.get("nickname") or c.get("username")
        print(f"    - {name} ({c.get('username', '')})")

results["contacts"] = {"count": r_cnt, "sample": r_contacts}

# ============================================================
# 模块4：群组
# ============================================================
print(f"\n{'='*50}")
print("模块4：群组")
print(f"{'='*50}")

r_groups = test("群组列表", lambda: rpc.get_groups({"limit": 50}))
if r_groups.get("success"):
    ct = r_groups.get("chatroom_table", [])
    rc = r_groups.get("rcontact_groups", [])
    print(f"    chatroom表: {len(ct)} 个")
    print(f"    rcontact群: {len(rc)} 个")
    for g in ct[:5]:
        name = g.get("chatroomnick") or g.get("chatroomname", "")
        mc = g.get("membercount", "?")
        print(f"    - {name} (成员:{mc})")

# 获取第一个群的成员
if r_groups.get("success") and r_groups.get("chatroom_table"):
    first_gid = r_groups["chatroom_table"][0].get("chatroomname", "")
    if first_gid:
        r_members = test(f"群成员({first_gid[:20]})", lambda: rpc.get_group_members({"group_id": first_gid}))
        if r_members.get("success"):
            print(f"    成员数: {r_members.get('count', 0)}")
            print(f"    前5成员: {r_members.get('member_ids', [])[:5]}")

results["groups"] = r_groups

# ============================================================
# 模块5：消息发送（WCDB直接写入）
# ============================================================
print(f"\n{'='*50}")
print("模块5：消息发送（WCDB直接写入）")
print(f"{'='*50}")

# 5a. 截图发送前
adb_sc(f"v3_before_send_{ts}")

# 5b. 发送消息
msg1 = f"[Frida WCDB验证 {datetime.now().strftime('%H:%M:%S')}] 好友:{r_cnt.get('type3_friends','?')}人 | 群:{r_cnt.get('groups','?')}个 | 直接DB写入"
r_send1 = test("发送消息到文件传输助手", lambda: rpc.send_message({"to_id": "filehelper", "content": msg1, "msg_type": "text"}))
print(f"    消息ID: {r_send1.get('message_id', 'N/A')}")
print(f"    方法: {r_send1.get('method', 'N/A')}")
time.sleep(1)

# 5c. 截图发送后
adb_sc(f"v3_after_send_{ts}")

# 5d. 验证DB中是否有消息
r_verify = test("验证消息已写入DB", lambda: rpc.get_messages({"talker": "filehelper", "limit": 5}))
if r_verify.get("success"):
    print(f"    DB中消息数: {r_verify.get('count', 0)}")
    for msg in r_verify.get("messages", [])[:3]:
        print(f"    - isSend={msg.get('isSend')} | {str(msg.get('content',''))[:60]}")

results["message_send"] = {"send": r_send1, "verify": r_verify}

# ============================================================
# 模块6：群发消息
# ============================================================
print(f"\n{'='*50}")
print("模块6：群发消息")
print(f"{'='*50}")

# 获取3个联系人ID用于测试
test_targets = []
if r_contacts.get("success"):
    for c in r_contacts.get("contacts", [])[:3]:
        test_targets.append(c.get("username", ""))
test_targets = [t for t in test_targets if t]

if test_targets:
    mass_content = f"[群发测试 {datetime.now().strftime('%H:%M:%S')}] Frida WCDB群发验证"
    r_mass = test(f"群发消息({len(test_targets)}人)", lambda: rpc.mass_message({"targets": test_targets, "content": mass_content}))
    if r_mass.get("success"):
        print(f"    成功: {r_mass.get('success_count')}/{r_mass.get('total')} 条")
    results["mass_message"] = r_mass
else:
    print("  ⚠️ 跳过群发（无联系人）")

# ============================================================
# 模块7：朋友圈
# ============================================================
print(f"\n{'='*50}")
print("模块7：朋友圈")
print(f"{'='*50}")

r_moments = test("读取朋友圈", lambda: rpc.get_moments({"limit": 10}))
if r_moments.get("success"):
    print(f"    朋友圈条数: {r_moments.get('count', 0)}")
    for m in r_moments.get("moments", [])[:3]:
        print(f"    - {str(m.get('content',''))[:60]}")

moment_content = f"[SDK验证{datetime.now().strftime('%m/%d %H:%M')}] Frida WCDB无感控制微信成功 好友:{r_cnt.get('type3_friends','?')}人"
r_post = test("发布朋友圈(WCDB)", lambda: rpc.post_moment({"content": moment_content}))
print(f"    SNS ID: {r_post.get('sns_id', 'N/A')}")

results["moments"] = {"read": r_moments, "post": r_post}

# ============================================================
# 模块8：标签管理
# ============================================================
print(f"\n{'='*50}")
print("模块8：标签管理")
print(f"{'='*50}")

r_labels = test("获取标签列表", rpc.get_labels)
if r_labels.get("success"):
    print(f"    标签数量: {r_labels.get('count', 0)}")
    for l in r_labels.get("labels", [])[:5]:
        print(f"    - [{l.get('labelId')}] {l.get('labelName')} (成员:{l.get('memberCount','?')})")

# 获取第一个标签的成员
if r_labels.get("success") and r_labels.get("labels"):
    first_label = r_labels["labels"][0]
    lid = first_label.get("labelId")
    if lid:
        r_lm = test(f"标签成员(ID={lid})", lambda: rpc.get_label_members({"label_id": lid}))
        if r_lm.get("success"):
            print(f"    成员数: {r_lm.get('count', 0)}")

results["labels"] = r_labels

# ============================================================
# 模块9：搜索
# ============================================================
print(f"\n{'='*50}")
print("模块9：搜索")
print(f"{'='*50}")

r_search1 = test("搜索联系人(关键词:游)", lambda: rpc.search_contacts({"keyword": "游"}))
if r_search1.get("success"):
    print(f"    结果: {r_search1.get('count', 0)} 条")
    for c in r_search1.get("results", [])[:3]:
        print(f"    - {c.get('nickname','')} ({c.get('username','')})")

r_search2 = test("搜索消息(关键词:验证)", lambda: rpc.search_messages({"keyword": "验证"}))
if r_search2.get("success"):
    print(f"    结果: {r_search2.get('count', 0)} 条")

results["search"] = {"contacts": r_search1, "messages": r_search2}

# ============================================================
# 模块10：好友操作
# ============================================================
print(f"\n{'='*50}")
print("模块10：好友操作")
print(f"{'='*50}")

r_remark = test("设置好友备注(filehelper)", lambda: rpc.set_remark({"wxid": "filehelper", "remark": "文件助手[SDK测试]"}))
print(f"    影响行数: {r_remark.get('affected_rows', 0)}")

results["friend_ops"] = {"remark": r_remark}

# ============================================================
# 模块11：原始SQL查询
# ============================================================
print(f"\n{'='*50}")
print("模块11：原始SQL查询（高级）")
print(f"{'='*50}")

r_sql1 = test("查询最近会话", lambda: rpc.raw_query({"sql": "SELECT username, content, createTime FROM conversation ORDER BY createTime DESC LIMIT 10"}))
if r_sql1.get("success"):
    print(f"    会话数: {r_sql1.get('count', 0)}")
    for row in r_sql1.get("rows", [])[:3]:
        print(f"    - {row.get('username','')} | {str(row.get('content',''))[:40]}")

r_sql2 = test("查询消息统计", lambda: rpc.raw_query({"sql": "SELECT talker, COUNT(*) as cnt FROM message GROUP BY talker ORDER BY cnt DESC LIMIT 10"}))
if r_sql2.get("success"):
    print(f"    活跃会话Top5:")
    for row in r_sql2.get("rows", [])[:5]:
        print(f"    - {row.get('talker','')} ({row.get('cnt',0)} 条)")

results["raw_sql"] = {"sessions": r_sql1, "stats": r_sql2}

# ============================================================
# 截图验证
# ============================================================
print(f"\n{'='*50}")
print("截图验证")
print(f"{'='*50}")

# 导航到文件传输助手截图
r_nav = rpc.navigate_to_chat({"wxid": "filehelper"})
print(f"  导航: {r_nav}")
sc = adb_sc(f"v3_final_chat_{ts}", wait=2)

# ============================================================
# 汇总报告
# ============================================================
total = passed + failed
pass_rate = (passed / total * 100) if total > 0 else 0

print(f"\n{'='*60}")
print("验证汇总报告")
print(f"{'='*60}")
print(f"  总测试项: {total}")
print(f"  通过: {passed} ✅")
print(f"  失败: {failed} ❌")
print(f"  通过率: {pass_rate:.1f}%")
print(f"\n  真实数据:")
print(f"    好友(type=3): {r_cnt.get('type3_friends', '?')} 人")
print(f"    总联系人: {r_cnt.get('total_rcontact', '?')} 条")
print(f"    群组: {r_cnt.get('groups', '?')} 个")
print(f"    标签: {r_labels.get('count', '?')} 个")
print(f"\n  消息发送: {'✅ WCDB直接写入' if r_send1.get('success') else '❌ 失败'}")
print(f"  群发消息: {'✅' if results.get('mass_message', {}).get('success') else '❌'}")
print(f"  朋友圈发布: {'✅' if r_post.get('success') else '❌'}")
print(f"  捕获真实发送类: {len(events)} 个事件")

# 保存完整报告
report = {
    "title": "工作手机SDK - 微信全功能Frida验证报告 v3",
    "time": datetime.now().isoformat(),
    "device": PHONE_IP,
    "wechat_pid": WECHAT_PID,
    "summary": {
        "total": total,
        "passed": passed,
        "failed": failed,
        "pass_rate": f"{pass_rate:.1f}%"
    },
    "real_data": {
        "type3_friends": r_cnt.get("type3_friends"),
        "total_rcontact": r_cnt.get("total_rcontact"),
        "groups": r_cnt.get("groups"),
        "labels": r_labels.get("count")
    },
    "results": results,
    "events": events
}
with open(RESULT_FILE, "w", encoding="utf-8") as f:
    json.dump(report, f, ensure_ascii=False, indent=2, default=str)
print(f"\n完整报告已保存: {RESULT_FILE}")

script.unload()
session.detach()
print("完成！")