#!/usr/bin/env python3
"""
WeChat 全库解密导出工具
步骤：
  1. Frida hook 持续捕获 AES key
  2. 每个新 key 试匹配所有未解锁的数据库
  3. 匹配成功 → 立即用 sqlcipher_export() 导出为明文 SQLite
  4. 所有库解完 or 超时后退出

用法：
  python3 ~/Downloads/wechat_export_all.py

输出：
  ~/Documents/wechat_plain/  目录下各数据库的明文副本
"""

import sys, os, time, subprocess, threading, shutil

# ── 路径 ──────────────────────────────────────────────────────────────
sys.path.insert(0, "/Users/yibo/Library/Python/3.9/lib/python/site-packages")
import frida

DB_BASE   = "/Users/yibo/Library/Containers/com.tencent.xinWeChat/Data/Documents/xwechat_files/yibooo_ac07/db_storage"
OUT_DIR   = "/Users/yibo/Documents/wechat_plain"
KEY_CACHE = "/Users/yibo/Documents/wechat_plain/keys.txt"
SQLCIPHER = "/opt/homebrew/bin/sqlcipher"

# ── 所有数据库（相对 DB_BASE）──────────────────────────────────────────
ALL_DBS = [
    # 消息（按大小降序，优先处理）
    "message/message_1.db",
    "message/message_2.db",
    "message/message_3.db",
    "message/message_4.db",
    "message/message_5.db",
    "message/message_6.db",
    "message/message_7.db",
    "message/message_8.db",
    "message/message_9.db",
    "message/message_10.db",
    "message/message_0.db",
    "message/message_fts.db",
    "message/media_0.db",
    "message/media_1.db",
    "message/message_resource.db",
    "message/biz_message_0.db",
    "message/biz_message_1.db",
    # 联系人
    "contact/contact.db",
    "contact/contact_fts.db",
    # 会话
    "session/session.db",
    # 其他
    "sns/sns.db",
    "general/general.db",
    "favorite/favorite.db",
    "favorite/favorite_fts.db",
    "head_image/head_image.db",
    "hardlink/hardlink.db",
    "bizchat/bizchat.db",
    "emoticon/emoticon.db",
    "migrate/unspportmsg.db",
    "solitaire/solitaire.db",
]

os.makedirs(OUT_DIR, exist_ok=True)

# ── 已知 key（上次运行结果）────────────────────────────────────────────
known_keys = {}          # hex_key → set of verified db relpath
pending_dbs = list(ALL_DBS)   # 还没匹配到 key 的数据库
exported_dbs = set()

def load_key_cache():
    """从缓存文件加载已知 key"""
    if not os.path.exists(KEY_CACHE):
        return
    with open(KEY_CACHE) as f:
        for line in f:
            line = line.strip()
            if ':' in line:
                relpath, key = line.split(':', 1)
                known_keys[key] = known_keys.get(key, set())
                known_keys[key].add(relpath)
                if relpath in pending_dbs:
                    pending_dbs.remove(relpath)
    if known_keys:
        print(f"[cache] 加载 {len(known_keys)} 个已知 key，{len(pending_dbs)} 个库待解")

def save_key_cache():
    with open(KEY_CACHE, 'w') as f:
        for key, dbs in known_keys.items():
            for db in dbs:
                f.write(f"{db}:{key}\n")

# ── sqlcipher 验证 ──────────────────────────────────────────────────────
def try_open(key_hex, relpath):
    """返回 True 表示 key 能打开这个数据库"""
    db_path = os.path.join(DB_BASE, relpath)
    if not os.path.exists(db_path):
        return False
    r = subprocess.run(
        [SQLCIPHER, db_path],
        input=f"PRAGMA key = \"x'{key_hex}'\";\nSELECT count(*) FROM sqlite_master;\n",
        capture_output=True, text=True, timeout=15
    )
    lines = [l for l in r.stdout.split('\n') if l.strip() and l.strip() != 'ok']
    return bool(lines) and lines[0].strip().isdigit()

# ── 导出函数 ────────────────────────────────────────────────────────────
export_lock = threading.Lock()

def export_db(key_hex, relpath):
    """用 sqlcipher_export 把加密库导出为明文 SQLite"""
    src  = os.path.join(DB_BASE, relpath)
    name = relpath.replace('/', '_')
    dst  = os.path.join(OUT_DIR, name)
    tmp  = dst + '.tmp'

    # 跳过已导出
    if relpath in exported_dbs:
        return
    if os.path.exists(dst):
        sz = os.path.getsize(dst)
        if sz > 4096:
            print(f"  [skip] {name} 已存在 ({sz//1024//1024} MB)")
            exported_dbs.add(relpath)
            return

    src_sz = os.path.getsize(src)
    print(f"\n  [export] {name} ({src_sz//1024//1024} MB) → {dst}")
    t0 = time.time()

    sql = "\n".join([
        f"PRAGMA key = \"x'{key_hex}'\";",
        "PRAGMA cipher_page_size = 4096;",
        f"ATTACH DATABASE '{tmp}' AS plain KEY '';",
        "SELECT sqlcipher_export('plain');",
        "DETACH DATABASE plain;",
    ])

    r = subprocess.run(
        [SQLCIPHER, src],
        input=sql, capture_output=True, text=True, timeout=600
    )

    elapsed = time.time() - t0
    if r.returncode == 0 and os.path.exists(tmp) and os.path.getsize(tmp) > 4096:
        os.rename(tmp, dst)
        out_sz = os.path.getsize(dst)
        print(f"  [✓] {name}: {src_sz//1024//1024} MB → {out_sz//1024//1024} MB 明文  ({elapsed:.1f}s)")
        exported_dbs.add(relpath)
    else:
        if os.path.exists(tmp):
            os.remove(tmp)
        err = r.stderr.strip()[:200]
        print(f"  [✗] {name}: 导出失败  {err}")

def match_key(key_hex):
    """用新 key 尝试匹配所有待解数据库，匹配成功后后台导出"""
    global pending_dbs
    matched = []
    for relpath in list(pending_dbs):
        try:
            if try_open(key_hex, relpath):
                matched.append(relpath)
        except Exception:
            pass

    if matched:
        known_keys[key_hex] = known_keys.get(key_hex, set()) | set(matched)
        for relpath in matched:
            pending_dbs.remove(relpath)
        save_key_cache()
        print(f"\n[key] {key_hex[:16]}... → 匹配 {len(matched)} 个库: {matched}")
        # 后台并发导出
        for relpath in matched:
            t = threading.Thread(target=export_db, args=(key_hex, relpath), daemon=True)
            t.start()
    return matched

# ── 对已知 key 立即导出 ─────────────────────────────────────────────────
def export_known():
    threads = []
    for key_hex, dbs in list(known_keys.items()):
        for relpath in dbs:
            t = threading.Thread(target=export_db, args=(key_hex, relpath), daemon=True)
            t.start()
            threads.append(t)
    return threads

# ── Frida 主逻辑 ─────────────────────────────────────────────────────────
WECHAT_APP = '/Applications/微信.app/Contents/MacOS/WeChat'
WECHAT_BUNDLE = 'com.tencent.xinWeChat'

def make_on_msg():
    def on_msg(msg, data):
        if msg['type'] != 'send': return
        p = msg['payload']
        if p.get('type') == 'ready':
            hooks = []
            if p.get('cc1'): hooks.append('CCCryptorCreate')
            if p.get('cc2'): hooks.append('CCCryptorCreateWithMode')
            if p.get('cc3'): hooks.append('CCCrypt')
            print(f"[Frida] hooks active: {', '.join(hooks) if hooks else '(none!)'}")
        elif p.get('type') == 'key':
            key_hex = p['hex']
            via = p.get('via', '?')
            print(f"\n[key captured via {via}] {key_hex[:16]}...")
            match_key(key_hex)
    return on_msg

def run_frida_earlybird():
    """关掉微信，等用户手动打开后立即 attach（earlybird 模式）"""
    existing = subprocess.run(['pgrep', '-x', 'WeChat'],
                              capture_output=True, text=True).stdout.strip()
    if existing:
        print(f"[*] 关闭现有 WeChat (PID={existing.split()[0]})...")
        subprocess.run(['killall', 'WeChat'], capture_output=True)
        time.sleep(2)

    print("[*] 请手动打开微信.app（双击 Dock 图标）")
    print("    脚本会在进程出现时立即 attach...\n")

    pid = None
    for _ in range(120):
        out = subprocess.run(['pgrep', '-x', 'WeChat'],
                             capture_output=True, text=True).stdout.strip()
        if out:
            pid = int(out.split('\n')[0])
            break
        time.sleep(0.5)

    if not pid:
        print("[!] 2 分钟内未检测到 WeChat，退出")
        return None

    print(f"[*] WeChat PID={pid}，立即 attach...")
    js_code = open('/Users/yibo/Downloads/wechat_key_hunter.js').read()
    session = frida.attach(pid)
    script = session.create_script(js_code)
    script.on('message', make_on_msg())
    script.load()
    print("[!] Hook 已就绪，请在微信里操作：")
    print("    · 滚动浏览聊天列表（触发各消息分片）")
    print("    · 打开朋友圈、收藏、搜索、表情、公众号")
    print("    · 按 Ctrl+C 停止监听并等待导出完成\n")
    return session

# ── 主程序 ──────────────────────────────────────────────────────────────
if __name__ == '__main__':
    print("=" * 60)
    print("WeChat 全库解密导出工具")
    print(f"输出目录: {OUT_DIR}")
    print("=" * 60)

    load_key_cache()

    # 先导出已知 key 对应的数据库
    threads = export_known()

    # 启动 Frida 捕获新 key
    if pending_dbs:
        print(f"\n[*] 还有 {len(pending_dbs)} 个库需要 key:")
        for db in pending_dbs:
            print(f"      {db}")

        session = run_frida_earlybird()
        if session:
            try:
                while pending_dbs:
                    time.sleep(2)
                    remaining = len(pending_dbs)
                    if remaining:
                        print(f"\r[待解] {remaining} 个库还需 key: "
                              f"{', '.join(db.split('/')[-1] for db in pending_dbs[:3])}"
                              f"{'...' if remaining > 3 else ''}   ", end='', flush=True)
            except KeyboardInterrupt:
                print("\n\n[*] 停止监听")
            finally:
                session.detach()
    else:
        print("\n[*] 所有 key 已从缓存加载，直接开始导出")

    # 等待所有导出线程完成
    print("\n[*] 等待导出任务完成...")
    active = threading.active_count()
    while threading.active_count() > 1:
        done = len(exported_dbs)
        total = len(ALL_DBS) - len(pending_dbs)
        print(f"\r    已完成 {done}/{total} 个数据库...  ", end='', flush=True)
        time.sleep(2)

    print(f"\n\n{'='*60}")
    print(f"完成！导出 {len(exported_dbs)} 个数据库")
    print(f"输出目录: {OUT_DIR}")
    print(f"{'='*60}")

    # 统计输出文件大小
    total_size = 0
    for f in os.listdir(OUT_DIR):
        if f.endswith('.db'):
            sz = os.path.getsize(os.path.join(OUT_DIR, f))
            total_size += sz
            print(f"  {f:<45} {sz//1024//1024:>5} MB")
    print(f"\n  总计: {total_size//1024//1024} MB")