#!/usr/bin/env python3
# frida CLI 的 `-W pattern` 只处理第一个匹配的 spawn，之后同 pattern 的新进程一律
# resume 但不 attach。微信"我的作品"发布一次流程会连续冒出多个 appbrand* 子进程
# （选图用一个、真正宿主 WebView 用另一个），必须逐个都 hook 到，所以改用
# device.enable_spawn_gating() + on('spawn-added') 手动处理每一个匹配的 spawn。
#
# 用法: source .venv/bin/activate && python3 multi_spawn_hook.py
import sys
import time

import frida

PATTERN = "com.tencent.mm:appbrand"
# session.create_script() (raw Python bindings) doesn't auto-bundle
# frida-java-bridge the way `frida -l script.js` (CLI) does, and there's no
# require() available either — must load a pre-compiled bundle instead.
# Regenerate after editing hook_lite_transfer.js's logic:
#   npm install frida-java-bridge (in a scratch dir with an entry.js that
#   does `import Java from 'frida-java-bridge';` + the hook body)
#   frida-compile entry.js -o hook_lite_transfer_bundled.js -P gum
SCRIPT_PATH = "../wechat_send_internal/hook_lite_transfer_bundled.js"


def on_message(pid, message, data):
    if message["type"] == "send":
        print(f"[pid={pid}] {message['payload']}")
    elif message["type"] == "error":
        print(f"[pid={pid}] ERROR: {message.get('stack', message)}")


def main():
    device = frida.get_usb_device(timeout=10)
    with open(SCRIPT_PATH) as f:
        source = f.read()

    def on_spawn_added(spawn):
        identifier = spawn.identifier or ""
        pid = spawn.pid
        if PATTERN in identifier:
            print(f"[+] matched spawn pid={pid} identifier={identifier}, attaching...")
            try:
                session = device.attach(pid)
                script = session.create_script(source)
                script.on("message", lambda m, d: on_message(pid, m, d))
                script.load()
                print(f"[+] hooked pid={pid}")
            except Exception as e:
                print(f"[!] hook failed for pid={pid}: {e}")
        else:
            print(f"[ ] ignoring spawn pid={pid} identifier={identifier}")
        device.resume(pid)

    device.on("spawn-added", on_spawn_added)
    device.enable_spawn_gating()
    print(f"[*] spawn gating enabled, watching for '{PATTERN}' ... (Ctrl+C to stop)")
    try:
        while True:
            time.sleep(3600)
    except KeyboardInterrupt:
        print("stopping")


if __name__ == "__main__":
    main()