chuanqi/cq2_sdk/scripts/frida_watch_capture.py

"""Auto-attach watcher: polls for the target process, attaches frida
the moment it appears, runs the capture, and exits.

Use when launching the game and you want the capture to start
automatically as soon as cq2.exe is up.

Usage:
  uv run python scripts/frida_watch_capture.py
  uv run python scripts/frida_watch_capture.py --name 996box
  uv run python scripts/frida_watch_capture.py --interval 2 --duration 60
"""
import argparse
import os
import subprocess
import sys
import time
from pathlib import Path

sys.path.insert(0, os.path.join(os.path.dirname(__file__), os.pardir, "src"))


def get_pids_by_name(name: str) -> list[int]:
    out: list[int] = []
    try:
        # PowerShell: Get-Process | Where-Object Name -like "*<name>*"
        ps_cmd = (
            f"Get-Process | Where-Object {{ $_.Name -like '*{name}*' }} "
            f"| Select-Object -ExpandProperty Id"
        )
        result = subprocess.run(
            ["powershell", "-NoProfile", "-Command", ps_cmd],
            capture_output=True, text=True, timeout=10,
        )
        for line in result.stdout.strip().splitlines():
            line = line.strip()
            if line.isdigit():
                out.append(int(line))
    except Exception as e:
        print(f"poll err: {e}", file=sys.stderr)
    return out


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--name", default="cq2", help="process name fragment to watch for")
    p.add_argument("--interval", type=float, default=2.0, help="poll interval (s)")
    p.add_argument("--timeout", type=float, default=600.0, help="max wait time (s)")
    p.add_argument("--duration", type=float, default=60.0, help="capture duration after attach (s)")
    p.add_argument("--out", default="cq2_capture.jsonl")
    p.add_argument("--frida-host", default=os.environ.get("FRIDA_HOST", "127.0.0.1:27042"))
    args = p.parse_args()

    print(f"watching for process matching '*{args.name}*' (timeout {args.timeout}s)")
    deadline = time.time() + args.timeout
    seen_pid: int | None = None
    while time.time() < deadline:
        pids = get_pids_by_name(args.name)
        if pids:
            seen_pid = pids[0]
            print(f"found pid={seen_pid}, launching capture...")
            break
        time.sleep(args.interval)

    if seen_pid is None:
        print(f"timed out waiting for '{args.name}'", file=sys.stderr)
        return 1

    # delegate to frida_capture.py with this pid
    cmd = [
        sys.executable, "-m", "uv", "run", "python",
        str(Path(__file__).parent / "frida_capture.py"),
        "--pid", str(seen_pid),
        "--out", args.out,
        "--duration", str(args.duration),
        "--frida-host", args.frida_host,
    ]
    print(f"running: {' '.join(cmd)}")
    return subprocess.call(cmd)


if __name__ == "__main__":
    sys.exit(main())