#!/usr/bin/env python3
"""
WeChat Android Full Toolkit v2
- Key capture from libWCDB.so
- DB extraction via adb
- Auto decryption
- Message monitoring (send/receive)
- Active message sending via Frida RPC
- ContentProvider auth bypass

Usage:
    python3 wechat_android_attach_v2.py              # Capture keys + monitor
    python3 wechat_android_attach_v2.py --extract-db # Pull DB from phone
    python3 wechat_android_attach_v2.py --decrypt    # Decrypt with captured key
    python3 wechat_android_attach_v2.py --send-msg <toUser> <content>  # Send message
    python3 wechat_android_attach_v2.py --status     # Check Frida session status
"""
import sys, os, json, time, subprocess, re

# ── Configuration ─────────────────────────────────────────────────────────────
WECHAT_PKG = "com.tencent.mm"
JS_PATH = os.path.join(os.path.dirname(__file__), "wechat_android_full_hunter.js")
OUT_DIR = os.path.expanduser("~/wxmoment")
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 = {}
cipher_specs = {}

# 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():
    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 get_wechat_pid(device):
    """Get WeChat main process PID via pgrep fallback."""
    # Try with root first
    for su_path in ["/debug_ramdisk/su", "/data/local/tmp/su2"]:
        try:
            result = subprocess.run(
                ["adb", "shell", su_path, "-c", "pgrep -x com.tencent.mm"],
                capture_output=True, text=True
            )
            if result.returncode == 0 and result.stdout.strip():
                pids = sorted([int(p.strip()) for p in result.stdout.strip().split("\n") if p.strip()])
                if pids:
                    return pids[0]  # Main process is the smallest PID
        except Exception:
            pass

    # Fallback without root
    try:
        result = subprocess.run(
            ["adb", "shell", "pgrep", "-x", "com.tencent.mm"],
            capture_output=True, text=True
        )
        pids = sorted([int(p.strip()) for p in result.stdout.strip().split("\n") if p.strip()])
        if pids:
            return pids[0]
    except Exception as e:
        print(f"[!] pgrep failed: {e}")

    try:
        for proc in device.enumerate_processes():
            if proc.name == WECHAT_PKG:
                return proc.pid
    except Exception as e:
        print(f"[!] enumerate_processes failed: {e}")

    return None

def attach_wechat(device, wait_for_process=True):
    """Attach to WeChat main process."""
    pid = get_wechat_pid(device)

    # Wait for process to appear if requested
    if not pid and wait_for_process:
        print("[*] Waiting for WeChat process to appear...")
        for _ in range(30):  # Wait up to 30 seconds
            time.sleep(1)
            pid = get_wechat_pid(device)
            if pid:
                print(f"[*] WeChat appeared (PID={pid})")
                break

    if pid:
        print(f"[*] Attaching to {WECHAT_PKG} (PID={pid})")
        try:
            return device.attach(pid)
        except Exception as e:
            print(f"[!] Attach to PID={pid} failed: {e}")

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

    return None

# ── Session management ────────────────────────────────────────────────────────
active_session = None
active_script = None

def load_script(session, js_path):
    global active_script
    with open(js_path) as f:
        js_code = f.read()
    script = session.create_script(js_code)
    script.on("message", on_message)
    script.load()
    active_script = script
    return script

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", {})

        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}")

    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 == "msg":
        direction = payload.get("direction", "")
        toUser = payload.get("toUser", "")
        content = payload.get("content", "")
        msgType = payload.get("msgType", 0)
        print(f"\n  [MSG {direction.upper()}] to={toUser} type={msgType}")
        print(f"    content: {content[:100]}{'...' if len(content) > 100 else ''}")

    elif ptype == "sql_insert":
        table = payload.get("table", "")
        sql = payload.get("sql", "")
        print(f"\n  [SQL INSERT] table={table}")
        print(f"    sql: {sql[:200]}{'...' if len(sql) > 200 else ''}")

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

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

# ── DB Extraction ─────────────────────────────────────────────────────────────
def get_su_cmd():
    for su_path in ["/debug_ramdisk/su", "/data/local/tmp/su2"]:
        try:
            result = subprocess.run(
                ["adb", "shell", su_path, "-c", "id"],
                capture_output=True, text=True
            )
            if result.returncode == 0 and "uid=0" in result.stdout:
                return ["adb", "shell", su_path, "-c"]
        except Exception:
            pass
    return ["adb", "shell", "su", "-c"]

def find_enmicromsg_db():
    print("[*] Scanning for EnMicroMsg.db on phone...")
    su_cmd = get_su_cmd()
    result = subprocess.run(
        su_cmd + ["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]

def extract_db(db_path=None, out_dir=DB_OUT_DIR):
    if db_path is None:
        db_path = find_enmicromsg_db()
    if not db_path:
        return False

    files_to_copy = [db_path, db_path + "-shm", db_path + "-wal"]
    os.makedirs(out_dir, exist_ok=True)
    su_cmd = get_su_cmd()

    for src in files_to_copy:
        dst = os.path.join(out_dir, os.path.basename(src))
        print(f"[*] Copying {src} -> {dst}")
        r = subprocess.run(
            su_cmd + [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

# ── Decryption ────────────────────────────────────────────────────────────────
def try_decrypt(db_path, key_hex, out_path=None):
    if out_path is None:
        out_path = db_path + ".decrypted.db"

    if os.path.exists(out_path):
        os.remove(out_path)

    try:
        key_str = bytes.fromhex(key_hex).decode('latin-1')
    except ValueError:
        key_str = key_hex

    candidates = []
    candidates.append(("string", f'PRAGMA key = "{key_str}";'))
    candidates.append(("hex", f"PRAGMA key = \"x'{key_hex}'\";"))

    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

# ── Message Sending ───────────────────────────────────────────────────────────
def send_text_msg(script, to_user, content, method="auto"):
    """Send text message via Frida RPC."""
    print(f"[*] Sending message to {to_user}...")
    print(f"    Content: {content[:50]}{'...' if len(content) > 50 else ''}")

    rpc = script.exports

    # Check status first
    try:
        status = rpc.get_status()
        print(f"    Status: pid={status.get('pid')}, javaAvailable={status.get('javaAvailable')}, javaHooksActive={status.get('javaHooksActive')}")
        if not status.get("javaAvailable"):
            print("[!] Java not available in target process. Cannot send messages via Java API.")
            print("    Try: spawn a fresh process, or attach to a different WeChat sub-process.")
            return False
    except Exception as e:
        print(f"[!] Failed to get status: {e}")
        return False

    # Try sending
    if method == "auto" or method == "rpc":
        try:
            result = rpc.send_text_msg(to_user, content)
            print(f"    [RPC Result] {json.dumps(result, ensure_ascii=False, indent=2)}")
            if result.get("success"):
                return True
        except Exception as e:
            print(f"    [RPC Error] {e}")

    if method == "auto" or method == "provider":
        try:
            result = rpc.send_via_provider(to_user, content)
            print(f"    [Provider Result] {json.dumps(result, ensure_ascii=False, indent=2)}")
            if result.get("success"):
                return True
        except Exception as e:
            print(f"    [Provider Error] {e}")

    return False

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

    print("=" * 60)
    print("WeChat Android Full Hunter v2")
    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 = attach_wechat(device)
    if not session:
        print("[!] Failed to attach to WeChat")
        sys.exit(1)

    global active_session
    active_session = session

    script = load_script(session, JS_PATH)

    print("[*] Script loaded. Hooks active.")
    print("[*] Operations:")
    print("    - Open WeChat, browse chats to trigger DB access")
    print("    - Send messages to monitor outgoing traffic")
    print("    - Press Ctrl+C to stop")
    print()

    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']}")

def run_send_msg(to_user, content):
    print("=" * 60)
    print("WeChat Send Message via Frida")
    print("=" * 60)

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

    device = get_usb_device()
    session = attach_wechat(device)
    if not session:
        print("[!] Failed to attach to WeChat")
        sys.exit(1)

    script = load_script(session, JS_PATH)
    time.sleep(1)  # Let hooks initialize

    success = send_text_msg(script, to_user, content)

    session.detach()
    return success

def run_status():
    device = get_usb_device()
    session = attach_wechat(device)
    if not session:
        print("[!] Failed to attach")
        sys.exit(1)

    script = load_script(session, JS_PATH)
    time.sleep(0.5)

    try:
        status = script.exports.get_status()
        print(json.dumps(status, indent=2))
    except Exception as e:
        print(f"[!] Error: {e}")

    session.detach()

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:
            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)

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

    if len(sys.argv) > 1 and sys.argv[1] == "--send-msg":
        if len(sys.argv) < 4:
            print("Usage: python3 wechat_android_attach_v2.py --send-msg <toUser> <content>")
            print("  toUser: wxid_xxx or xxx@chatroom")
            print("  content: message text")
            sys.exit(1)
        success = run_send_msg(sys.argv[2], sys.argv[3])
        sys.exit(0 if success else 1)

    if len(sys.argv) > 1 and sys.argv[1] == "--status":
        run_status()
        return

    run_capture_and_monitor()

if __name__ == "__main__":
    main()