#!/usr/bin/env python3
"""
BlooMoo Script Engine Tracer  -  Frida launcher
Spawns nemo.exe and instruments bloomoodll.dll.

Usage:
    python frida_launcher.py                       # spawn nemo.exe
    python frida_launcher.py --attach <PID>        # attach to running process
    python frida_launcher.py --exe path/nemo.exe   # custom exe path
    python frida_launcher.py --script bloomoo_event_trace.js

Alternatively, use the Frida CLI directly:
    frida -l bloomoo_trace.js -f nemo.exe
"""

import argparse
import os
import sys
import threading
import shlex

try:
    import frida
except ImportError:
    print("Frida not installed. Run:  pip install frida-tools")
    sys.exit(1)


def format_trace(message):
    """Extract printable text from a Frida message. Returns (text, is_trace)."""
    mtype = message.get("type", "")
    if mtype == "send":
        payload = message.get("payload", "")
        if isinstance(payload, dict) and payload.get("type") == "trace":
            return payload.get("text", ""), True
        return f"[send] {payload}", False
    if mtype == "error":
        desc = message.get("description", "")
        stack = message.get("stack", "")
        return f"[error] {desc}\n{stack}" if stack else f"[error] {desc}", False
    return f"[{mtype}] {message}", False


def main():
    parser = argparse.ArgumentParser(
        description="BlooMoo Script Engine Tracer"
    )
    parser.add_argument(
        "--exe", default="nemo.exe",
        help="Path to nemo.exe (default: nemo.exe in script dir)"
    )
    parser.add_argument(
        "--script", default="bloomoo_trace.js",
        help="Frida JS script to load (default: bloomoo_trace.js)"
    )
    parser.add_argument(
        "--attach", type=int, default=0, metavar="PID",
        help="Attach to an already-running process instead of spawning"
    )
    parser.add_argument(
        "--log", default="bloomoo_trace.log", metavar="FILE",
        help="Write trace output to a file (default: bloomoo_trace.log, use --no-log to disable)"
    )
    parser.add_argument(
        "--no-log", action="store_true",
        help="Disable file logging"
    )
    parser.add_argument(
        "--args", default="-window", metavar="ARGS",
        help="Additional arguments to pass to the spawned process"
    )
    args = parser.parse_args()

    script_dir = os.path.dirname(os.path.abspath(__file__))
    js_path = args.script
    if not os.path.isabs(js_path):
        js_path = os.path.join(script_dir, js_path)

    if not os.path.isfile(js_path):
        print(f"[!] Script not found: {js_path}")
        sys.exit(1)

    with open(js_path, "r", encoding="utf-8") as f:
        js_code = f.read()

    # log file (on by default)
    log_file = None
    if args.log and not args.no_log:
        log_path = os.path.join(script_dir, args.log) if not os.path.isabs(args.log) else args.log
        log_file = open(log_path, "w", encoding="utf-8")
        print(f"[*] Logging to {log_path}")

    def on_msg_wrapper(message, data):
        text, is_trace = format_trace(message)
        if not text:
            return
        # trace lines already go to stdout via Frida's console.log;
        # non-trace messages (errors etc.) we print ourselves
        if not is_trace:
            print(text, flush=True)
        # write ALL messages (trace + errors) to the log file
        if log_file:
            try:
                log_file.write(text + "\n")
                log_file.flush()
            except Exception:
                pass

    device = frida.get_local_device()

    if args.attach:
        print(f"[*] Attaching to PID {args.attach} ...")
        session = device.attach(args.attach)
    else:
        exe_path = args.exe
        if not os.path.isabs(exe_path):
            exe_path = os.path.join(script_dir, exe_path)
        if not os.path.isfile(exe_path):
            print(f"[!] Executable not found: {exe_path}")
            sys.exit(1)
        print(f"[*] Spawning {exe_path} with args {args.args} ...")
        pid = device.spawn([exe_path] + shlex.split(args.args), cwd=script_dir)
        print(f"[*] PID = {pid}")
        session = device.attach(pid)

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

    if not args.attach:
        print("[*] Resuming process ...")
        device.resume(pid)

    print("[*] Tracer running.  Press Ctrl+C to detach and exit.")
    print("[*] Tip: hookGetObject / hookGetVariable are OFF by default (noisy).")
    print()

    try:
        # keep the script alive
        sys.stdin.read()
    except KeyboardInterrupt:
        pass

    print("\n[*] Detaching ...")
    session.detach()
    if log_file:
        log_file.close()
    print("[*] Done.")


if __name__ == "__main__":
    main()