"""
WeChat early-bird key capture:
1. Kill WeChat
2. Poll until new WeChat process appears
3. Attach Frida immediately
4. Capture CCCryptorCreate keys during login / DB open
"""
import sys, os, time, subprocess, threading
sys.path.insert(0, '/Users/yibo/Library/Python/3.9/lib/python/site-packages')
import frida

SQLCIPHER = '/opt/homebrew/bin/sqlcipher'
DB_BASE   = '/Users/yibo/Library/Containers/com.tencent.xinWeChat/Data/Documents/xwechat_files/yibooo_ac07/db_storage'

JS = r"""
var seen = {};

function uniqueByteCount(arr) {
    var map = {};
    var n = 0;
    for (var i = 0; i < arr.length; i++) {
        if (!map[arr[i]]) { map[arr[i]] = true; n++; }
    }
    return n;
}

function tryCapture(keyPtr, keyLen, label) {
    try {
        if (keyLen !== 32) return;
        var raw = keyPtr.readByteArray(32);
        if (!raw) return;
        var bytes = new Uint8Array(raw);
        if (uniqueByteCount(bytes) < 8) return;
        var hex = '';
        for (var i = 0; i < bytes.length; i++) {
            hex += ('0' + bytes[i].toString(16)).slice(-2);
        }
        if (!seen[hex]) {
            seen[hex] = true;
            send({type:'key', hex:hex, len:keyLen, via:label});
        }
    } catch(e) {
        send({type:'err', msg: e.toString(), label: label});
    }
}

var found = {};

var addr1 = Module.findGlobalExportByName('CCCryptorCreate');
found['CCCryptorCreate'] = !!addr1;
if (addr1) {
    Interceptor.attach(addr1, { onEnter: function(args) {
        if (args[1].toInt32() === 0) tryCapture(args[3], args[4].toInt32(), 'CCCryptorCreate');
    }});
}

var addr2 = Module.findGlobalExportByName('CCCryptorCreateWithMode');
found['CCCryptorCreateWithMode'] = !!addr2;
if (addr2) {
    Interceptor.attach(addr2, { onEnter: function(args) {
        if (args[2].toInt32() === 0) tryCapture(args[5], args[6].toInt32(), 'CCCryptorCreateWithMode');
    }});
}

var addr3 = Module.findGlobalExportByName('CCCrypt');
found['CCCrypt'] = !!addr3;
if (addr3) {
    Interceptor.attach(addr3, { onEnter: function(args) {
        if (args[1].toInt32() === 0) tryCapture(args[3], args[4].toInt32(), 'CCCrypt');
    }});
}

send({type:'ready', found:found});
"""

def verify_key(key_hex, rel):
    path = os.path.join(DB_BASE, rel)
    if not os.path.exists(path): return False
    r = subprocess.run([SQLCIPHER, path],
        input=f"PRAGMA key=\"x'{key_hex}'\";\nSELECT count(*) FROM sqlite_master;\n",
        capture_output=True, text=True, timeout=10)
    lines = [l.strip() for l in r.stdout.split('\n') if l.strip() and l.strip() != 'ok']
    return bool(lines) and lines[0].isdigit()

TEST_DBS = [
    'contact/contact.db',
    'message/message_0.db',
    'session/session.db',
    'general/general.db',
    'message/message_1.db',
]

confirmed = {}  # key_hex -> [dbs]

def on_key(key_hex, via, keylen):
    print(f'\n  [key] len={keylen} via={via}  {key_hex[:24]}...')
    matched = []
    for rel in TEST_DBS:
        try:
            if verify_key(key_hex, rel):
                matched.append(rel)
                print(f'    ✓ matches {rel}')
        except Exception as e:
            pass
    if not matched:
        print(f'    (no match in test set)')
    else:
        confirmed[key_hex] = matched

# ── main ──────────────────────────────────────────────────────────────────────
print('[1] Killing WeChat...')
subprocess.run(['killall', 'WeChat'], capture_output=True)
time.sleep(2)

print('[2] Waiting for WeChat to start (open 微信.app now)...')
print('    → 请现在手动打开微信.app')

pid = None
for _ in range(120):   # wait up to 2 min
    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('[!] WeChat did not start within 2 minutes')
    sys.exit(1)

print(f'[3] WeChat PID={pid} detected, attaching Frida NOW...')
session = frida.attach(pid)
script  = session.create_script(JS)

def on_msg(msg, data):
    if msg['type'] == 'error':
        print(f'[JS ERROR] {msg["description"]}')
        return
    if msg['type'] != 'send': return
    p = msg['payload']
    if p['type'] == 'ready':
        print(f'[4] Hooks: {p["found"]}')
        print('    → 请登录微信，然后慢慢浏览各个聊天窗口\n')
    elif p['type'] == 'key':
        on_key(p['hex'], p['via'], p['len'])
    elif p['type'] == 'err':
        print(f'  [js-err] {p["label"]}: {p["msg"]}')

script.on('message', on_msg)
script.load()
print('[*] Script loaded. Monitoring... (Ctrl+C to stop)\n')

try:
    while True:
        time.sleep(1)
except KeyboardInterrupt:
    pass

print(f'\n[结果] 捕获到 {len(confirmed)} 个已验证 key:')
for k, dbs in confirmed.items():
    print(f'  {k}  →  {dbs}')
session.detach()