#!/usr/bin/env python3
"""
WeChat Android SQLCipher Key Capture + DB Extraction
Runs in Termux / Debian proot, connects to frida-server on Android host.

Usage:
    python3 wechat_android_attach.py              # Capture keys
    python3 wechat_android_attach.py --extract-db # Pull DB from phone
    python3 wechat_android_attach.py --decrypt    # Decrypt with captured key

Prerequisites:
    - frida-server running on Android (as root)
    - frida Python package installed in proot/Termux
    - WeChat installed and logged in
    - Phone has root access
"""
import sys, os, json, time, subprocess, re

# ── Configuration ─────────────────────────────────────────────────────────────
WECHAT_PKG = "com.tencent.mm"
JS_PATH = os.path.join(os.path.dirname(__file__), "wechat_android_key_hunter.js")
OUT_DIR = "/root/wxmoment"  # inside Debian proot
KEY_CACHE = os.path.join(OUT_DIR, "wechat_keys.json")
DB_OUT_DIR = os.path.join(OUT_DIR, "dbs")

os.makedirs(OUT_DIR, exist_ok=True)
os.makedirs(DB_OUT_DIR, exist_ok=True)

# ── Key storage ───────────────────────────────────────────────────────────────
captured_keys = {}      # hex_key -> {path, via, timestamp, ...}
cipher_specs = {}       # field -> value

# Try loading existing keys
def load_keys():
    if not os.path.exists(KEY_CACHE):
        return
    try:
        with open(KEY_CACHE) as f:
            data = json.load(f)
            captured_keys.update(data.get("keys", {}))
            cipher_specs.update(data.get("cipher_specs", {}))
        print(f"[cache] Loaded {len(captured_keys)} known keys")
    except Exception as e:
        print(f"[!] Failed to load key cache: {e}")

def save_keys():
    try:
        with open(KEY_CACHE, "w") as f:
            json.dump({
                "keys": captured_keys,
                "cipher_specs": cipher_specs,
                "updated": time.time()
            }, f, indent=2, ensure_ascii=False)
    except Exception as e:
        print(f"[!] Failed to save keys: {e}")

# ── Frida setup ───────────────────────────────────────────────────────────────
try:
    import frida
except ImportError:
    print("[!] frida Python package not installed.")
    print("    Install with: pip install frida")
    sys.exit(1)

def get_usb_device():
    """Get USB-connected Frida device (Android phone)."""
    for _ in range(5):
        try:
            device = frida.get_usb_device(timeout=3)
            return device
        except Exception as e:
            print(f"[!] Waiting for USB device... ({e})")
            time.sleep(1)
    raise RuntimeError("No USB Frida device found. Is frida-server running?")

def spawn_or_attach(device, pkg):
    """Try to attach to running process, or spawn a new one."""
    # Check if already running
    try:
        apps = device.enumerate_applications()
        for app in apps:
            if app.identifier == pkg:
                if app.pid != 0:
                    print(f"[*] Attaching to running {pkg} (PID={app.pid})")
                    return device.attach(app.pid)
                break
    except Exception as e:
        print(f"[!] enumerate_applications failed: {e}")

    # Try direct attach
    try:
        session = device.attach(pkg)
        print(f"[*] Attached to {pkg}")
        return session
    except Exception as e:
        print(f"[!] Direct attach failed: {e}")

    # Spawn new
    print(f"[*] Spawning {pkg}...")
    pid = device.spawn([pkg])
    session = device.attach(pid)
    device.resume(pid)
    print(f"[*] Spawned and attached (PID={pid})")
    return session

# ── Message handler ───────────────────────────────────────────────────────────
def on_message(msg, data):
    if msg.get("type") != "send":
        return
    payload = msg.get("payload", {})
    ptype = payload.get("type", "")

    if ptype == "key":
        key_hex = payload.get("hex", "")
        key_len = payload.get("len", 0)
        via = payload.get("via", "")
        path = payload.get("path", "")
        extra = payload.get("extra", {})

        # Deduplicate
        cache_key = key_hex + "@" + path
        if cache_key in captured_keys:
            return

        captured_keys[cache_key] = {
            "hex": key_hex,
            "len": key_len,
            "via": via,
            "path": path,
            "extra": extra,
            "ts": time.time(),
            "ts_fmt": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        save_keys()

        print(f"\n  [KEY CAPTURED via {via}]")
        print(f"    hex: {key_hex[:32]}{'...' if len(key_hex) > 32 else ''}")
        print(f"    len: {key_len}")
        if path:
            print(f"    path: {path}")
        if extra:
            print(f"    extra: {json.dumps(extra, ensure_ascii=False)}")

        # If this is EnMicroMsg.db, suggest extraction
        if "EnMicroMsg" in path or not path:
            print(f"    => Run: python3 wechat_android_attach.py --extract-db")

    elif ptype == "cipher_spec":
        field = payload.get("field", "")
        value = payload.get("value", "")
        cipher_specs[field] = value
        save_keys()
        print(f"  [CIPHER SPEC] {field} = {value}")

    elif ptype == "info":
        info_msg = payload.get("msg", "")
        # Highlight DB path discoveries
        if info_msg.startswith("FOUND_DB:"):
            print(f"\n  [DB FOUND] {info_msg}")
        else:
            print(f"  [INFO] {info_msg}")

    elif ptype == "error":
        print(f"  [ERROR] {payload.get('msg', '')}")

# ── DB Extraction ─────────────────────────────────────────────────────────────
def find_enmicromsg_db():
    """Find EnMicroMsg.db path on phone via adb."""
    print("[*] Scanning for EnMicroMsg.db on phone...")
    result = subprocess.run(
        ["adb", "shell", "su", "-c",
         "find /data/data/com.tencent.mm/MicroMsg -name 'EnMicroMsg.db' -type f 2>/dev/null"],
        capture_output=True, text=True
    )
    paths = [p.strip() for p in result.stdout.strip().split("\n") if p.strip()]
    if not paths:
        print("[!] EnMicroMsg.db not found. Is WeChat logged in?")
        return None
    print(f"[+] Found {len(paths)} EnMicroMsg.db:")
    for i, p in enumerate(paths):
        print(f"    [{i}] {p}")
    return paths[0]  # Return first (usually the active account)

def extract_db(db_path=None, out_dir=DB_OUT_DIR):
    """Copy WeChat DB files from phone to local directory via adb shell su -c cat."""
    if db_path is None:
        db_path = find_enmicromsg_db()
    if not db_path:
        return False

    # Also copy -shm and -wal files
    files_to_copy = [db_path, db_path + "-shm", db_path + "-wal"]

    os.makedirs(out_dir, exist_ok=True)

    for src in files_to_copy:
        dst = os.path.join(out_dir, os.path.basename(src))
        print(f"[*] Copying {src} -> {dst}")
        r = subprocess.run(
            ["adb", "shell", "su", "-c", f"cat '{src}'"],
            capture_output=True
        )
        if r.returncode == 0 and r.stdout:
            with open(dst, "wb") as f:
                f.write(r.stdout)
            print(f"    [OK] {dst} ({os.path.getsize(dst)} bytes)")
        else:
            print(f"    [SKIP] {src} (not found or unreadable)")

    print(f"\n[+] DB files in: {out_dir}")
    return True

def try_decrypt(db_path, key_hex, out_path=None):
    """Try to decrypt EnMicroMsg.db with captured key using sqlcipher.
    Automatically tries string password and raw hex key formats,
    with common cipher configurations."""
    if out_path is None:
        out_path = db_path + ".decrypted.db"

    # Remove existing output
    if os.path.exists(out_path):
        os.remove(out_path)

    # Build candidate keys and configs
    try:
        key_str = bytes.fromhex(key_hex).decode('latin-1')
    except ValueError:
        key_str = key_hex

    candidates = []
    # String password format (most common for WCDB)
    candidates.append(("string", f'PRAGMA key = "{key_str}";'))
    # Raw hex key format
    candidates.append(("hex", f"PRAGMA key = \"x'{key_hex}'\";"))

    # Cipher configs to try
    page_size = cipher_specs.get("pageSize", 1024)
    kdf_iter = cipher_specs.get("kdfIteration", None)

    configs = []
    base = f"PRAGMA cipher_compatibility = 1;\nPRAGMA cipher_page_size = {page_size};"
    configs.append(("compat1", base))
    if kdf_iter:
        configs.append(("compat1+kdf", f"{base}\nPRAGMA kdf_iter = {kdf_iter};"))
    configs.append(("no_compat", f"PRAGMA cipher_page_size = {page_size};"))

    print(f"[*] Trying decryption of {os.path.basename(db_path)}...")
    print(f"    Key candidates: {len(candidates)} | Config candidates: {len(configs)}")

    for key_name, key_cmd in candidates:
        for cfg_name, cfg_pragmas in configs:
            sql = f"{key_cmd}\n{cfg_pragmas}\nATTACH DATABASE '{out_path}' AS plain KEY '';\nSELECT sqlcipher_export('plain');\nDETACH DATABASE plain;\n"

            r = subprocess.run(
                ["sqlcipher", db_path],
                input=sql, capture_output=True, text=True, timeout=300
            )

            if r.returncode == 0 and os.path.exists(out_path) and os.path.getsize(out_path) > 4096:
                sz = os.path.getsize(out_path)
                print(f"    [SUCCESS] Decrypted: {sz} bytes (key={key_name}, cfg={cfg_name})")
                print(f"    Use: sqlite3 {out_path}")
                return True

    err = r.stderr.strip()[:200] if r.stderr else "unknown"
    print(f"    [FAILED] All combinations failed. Last error: {err}")
    return False

# ── Main interactive loop ─────────────────────────────────────────────────────
def run_capture():
    load_keys()

    print("=" * 60)
    print("WeChat Android SQLCipher Key Capture")
    print("=" * 60)
    print(f"[*] JS script: {JS_PATH}")
    print(f"[*] Key cache: {KEY_CACHE}")
    print(f"[*] DB output: {DB_OUT_DIR}")
    print()

    if not os.path.exists(JS_PATH):
        print(f"[!] JS script not found: {JS_PATH}")
        sys.exit(1)

    device = get_usb_device()
    print(f"[*] Device: {device.name}")

    session = spawn_or_attach(device, WECHAT_PKG)

    with open(JS_PATH) as f:
        js_code = f.read()

    script = session.create_script(js_code)
    script.on("message", on_message)
    script.load()

    print("[*] Script loaded. Hooks active.")
    print("[*] Open WeChat, browse chats to trigger DB access.")
    print("[*] Press Ctrl+C to stop.\n")

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n[*] Stopping...")
    finally:
        session.detach()
        save_keys()
        print(f"\n[RESULT] Captured {len(captured_keys)} unique keys:")
        for k, v in captured_keys.items():
            print(f"  {v['hex'][:32]}... via={v['via']} path={v['path']}")

def main():
    if len(sys.argv) > 1 and sys.argv[1] == "--extract-db":
        extract_db()
        return

    if len(sys.argv) > 1 and sys.argv[1] == "--decrypt":
        db_path = sys.argv[2] if len(sys.argv) > 2 else None
        if not db_path:
            # Find most recent extracted DB
            dbs = [f for f in os.listdir(DB_OUT_DIR) if f.endswith(".db") and "decrypted" not in f]
            if not dbs:
                print("[!] No DB found. Run with --extract-db first.")
                sys.exit(1)
            db_path = os.path.join(DB_OUT_DIR, sorted(dbs, key=lambda x: os.path.getmtime(os.path.join(DB_OUT_DIR, x)))[-1])

        load_keys()
        if not captured_keys:
            print("[!] No keys captured yet. Run capture first.")
            sys.exit(1)

        # Try each key
        for k, v in captured_keys.items():
            if try_decrypt(db_path, v["hex"]):
                break
        return

    run_capture()

if __name__ == "__main__":
    main()