#!/usr/bin/env python3
"""
Frida bridge server: runs on Mac, accepts frida injection requests from Android device.
Connects to frida-server on device via USB and executes JS scripts.
Port: 8767 (mapped to device via `adb reverse tcp:8767 tcp:8767`)

Architecture (frida 17.x):
- Java/ObjC globals are lazy-loaded via frida:load-bridge message protocol
- Python side must intercept these messages and send back the bridge JS source
- We pre-load the REPL agent (repl_agent.js) which sets up the lazy getters
"""
import threading
import json
from pathlib import Path
from typing import Any, Mapping

import frida
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

app = FastAPI()

# Locate frida_tools bridges directory (contains java.js, objc.js, etc.)
try:
    import frida_tools
    BRIDGES_DIR = Path(frida_tools.__file__).parent / "bridges"
    REPL_AGENT_JS = (Path(frida_tools.__file__).parent / "repl_agent.js").read_text(encoding="utf-8")
except Exception:
    BRIDGES_DIR = None
    REPL_AGENT_JS = None


class FridaRequest(BaseModel):
    package: str
    script: str
    timeout: int = 10


def try_handle_bridge_request(message: Mapping[Any, Any], script: frida.core.Script) -> bool:
    """Handle frida 17.x lazy bridge loading protocol (frida:load-bridge / frida:bridge-loaded)."""
    if message.get("type") != "send":
        return False
    payload = message.get("payload")
    if not isinstance(payload, dict) or payload.get("type") != "frida:load-bridge":
        return False
    name = payload["name"].lower()
    if BRIDGES_DIR is None:
        return False
    bridge_path = BRIDGES_DIR / f"{name}.js"
    if not bridge_path.exists():
        return False
    script.post({
        "type": "frida:bridge-loaded",
        "filename": bridge_path.name,
        "source": bridge_path.read_text(encoding="utf-8"),
    })
    return True


def _build_bundled_script(user_script: str) -> str:
    """
    Wrap user script with the REPL agent bundle so the Java global is available.
    The REPL agent (repl_agent.js) registers lazy getters for Java/ObjC/Swift
    that trigger the frida:load-bridge protocol when first accessed.
    """
    if REPL_AGENT_JS is None or not REPL_AGENT_JS.startswith("📦\n"):
        return user_script

    # User script runs after the agent; wrap in eval to catch errors
    user_code = user_script.strip()
    user_size = len(user_code.encode("utf-8"))
    # Append as extra module in the bundle (frida will load all modules)
    bundled = REPL_AGENT_JS.rstrip("\n") + f"\n✄\n{user_size} /frida/user-script.js\n✄\n{user_code}"
    return bundled


@app.get("/health")
def health():
    return {
        "status": "ok",
        "frida_version": frida.__version__,
        "bridges_ok": BRIDGES_DIR is not None and BRIDGES_DIR.exists(),
        "repl_agent_ok": REPL_AGENT_JS is not None,
    }


@app.post("/frida_call")
def frida_call(req: FridaRequest):
    results = []
    done = threading.Event()
    error_info = {}
    script_holder = [None]

    def on_message(msg, data):
        # Handle frida 17.x Java bridge loading protocol
        if script_holder[0] and try_handle_bridge_request(msg, script_holder[0]):
            return
        t = msg.get("type")
        if t == "send":
            payload = msg.get("payload")
            # Filter out internal frida protocol messages
            if isinstance(payload, dict) and str(payload.get("type", "")).startswith("frida:"):
                return
            results.append(payload)
            done.set()
        elif t == "error":
            error_info["error"] = msg.get("description", str(msg))
            done.set()

    try:
        dev = frida.get_usb_device(timeout=5)
        procs = dev.enumerate_processes()
        # Prefer main process (no colon in name like :push, :appbrand)
        matching = [p for p in procs if req.package in p.name and ":" not in p.name]
        if not matching:
            matching = [p for p in procs if req.package in p.name]
        if not matching:
            return {"error": f"process not found: {req.package}"}
        target = sorted(matching, key=lambda p: p.pid)[0]

        session = dev.attach(target.pid)
        full_source = _build_bundled_script(req.script)
        script_obj = session.create_script(full_source)
        script_holder[0] = script_obj
        script_obj.on("message", on_message)
        script_obj.load()
        done.wait(timeout=req.timeout)
        try:
            script_obj.unload()
        except Exception:
            pass
        session.detach()

        if error_info:
            return {"error": error_info["error"], "partial": results}
        return {"output": results, "count": len(results), "pid": target.pid}
    except frida.ProcessNotFoundError:
        return {"error": f"process not found: {req.package}"}
    except Exception as e:
        return {"error": str(e)}


if __name__ == "__main__":
    print(f"Frida {frida.__version__} bridge server — port 8767")
    print(f"Bridges: {BRIDGES_DIR}")
    uvicorn.run(app, host="0.0.0.0", port=8767)