#!/usr/bin/env python3
import sys, os, time, hashlib, hmac, struct, subprocess

sys.path.insert(0, os.path.expanduser("~/Library/Python/3.9/lib/python/site-packages"))
import frida

DB_BASE = os.path.expanduser(
    "~/Library/Containers/com.tencent.xinWeChat/Data/Documents/"
    "xwechat_files/yibooo_ac07/db_storage/"
)

def load_header(relpath):
    p = os.path.join(DB_BASE, relpath)
    return open(p, 'rb').read(4096) if os.path.exists(p) else None

def verify_raw_v4(key_hex, header):
    try:
        key = bytes.fromhex(key_hex)
        salt = header[:16]
        mac_salt = bytes(b ^ 0x3A for b in salt)
        hk = hashlib.pbkdf2_hmac('sha256', key, mac_salt, 2, dklen=32)
        h = hmac.new(hk, digestmod='sha256')
        h.update(header[16:4048]); h.update(header[4048:4064]); h.update(struct.pack('<I', 1))
        return h.digest() == header[4064:4096]
    except: return False

def verify_pbkdf2_v4(key_hex, header):
    try:
        key = bytes.fromhex(key_hex)
        salt = header[:16]
        dk = hashlib.pbkdf2_hmac('sha256', key, salt, 256000, dklen=64)
        mac_key = dk[32:]
        mac_salt = bytes(b ^ 0x3A for b in salt)
        hk = hashlib.pbkdf2_hmac('sha256', mac_key, mac_salt, 2, dklen=32)
        h = hmac.new(hk, digestmod='sha256')
        h.update(header[16:4048]); h.update(header[4048:4064]); h.update(struct.pack('<I', 1))
        return h.digest() == header[4064:4096]
    except: return False

headers = {}
for name, relpath in [('contact','contact/contact.db'),('msg0','message/message_0.db'),('general','general/general.db')]:
    h = load_header(relpath)
    if h: headers[name] = h
print(f"已加载 {len(headers)} 个DB头部")

pid_out = subprocess.run(['pgrep','-x','WeChat'], capture_output=True, text=True).stdout.strip()
if not pid_out:
    print("微信未运行"); sys.exit(1)
pid = int(pid_out.split('\n')[0])
print(f"附加 PID={pid}...")

session = frida.attach(pid)
print("Frida attached!")

script = session.create_script(open('/tmp/wechat_hook.js').read())
captured = []

def on_message(msg, data):
    if msg['type'] != 'send': 
        if msg['type'] == 'error': print(f"[ScriptErr] {msg.get('description','')}")
        return
    p = msg['payload']
    t = p.get('type')
    if t == 'ready': print(f"[Frida] {p['msg']}")
    elif t == 'hooked': print(f"[Hook] {p['func']} @ {p['module']}")
    elif t == 'key':
        k = p['key_hex']; n = p['key_len']; fn = p['func']
        print(f"\n[KEY] {fn} len={n} hex={k}")
        try:
            printable = ''.join(chr(b) if 32<=b<127 else '.' for b in bytes.fromhex(k))
            print(f"      str={printable}")
        except: pass
        for dbname, hdr in headers.items():
            if verify_raw_v4(k, hdr):
                print(f"      ✅ raw_v4 OK: {dbname}")
                captured.append(('raw_v4', k, dbname))
                open('/tmp/wechat_key.hex','w').write(k)
            elif verify_pbkdf2_v4(k, hdr):
                print(f"      ✅ pbkdf2_v4 OK: {dbname}")
                captured.append(('pbkdf2_v4', k, dbname))
                open('/tmp/wechat_key.hex','w').write(k)
    elif t == 'error': print(f"[Err] {p['msg']}")

script.on('message', on_message)
script.load()

print("\n监听 sqlite3_key 调用... 请在微信里操作触发DB访问")
print("Ctrl+C 退出\n")
try:
    while True: time.sleep(1)
except KeyboardInterrupt: pass

session.detach()
if captured:
    for mode, k, db in captured:
        print(f"[{mode}] key={k} ({db})")
else:
    print("未捕获到key（DB已在attach前打开，需重启微信后立即运行）")