#!/usr/bin/env python3
"""Orchestrate an automated PE binary exploration workflow via peek-a-bin's MCP server.

Loads a PE file, runs anomaly detection, enumerates and filters functions,
decompiles the entry point and suspicious functions, collects cross-references,
and produces a structured summary report.  Communication with the MCP server
happens over stdin/stdout using JSON-RPC messages.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import textwrap
from pathlib import Path
from typing import Any


# ---------------------------------------------------------------------------
# MCP client helpers
# ---------------------------------------------------------------------------

class McpClient:
    """Minimal MCP client that talks to a peek-a-bin MCP server over stdio."""

    def __init__(self, mcp_command: list[str], cwd: str | None = None) -> None:
        self._proc = subprocess.Popen(
            mcp_command,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            cwd=cwd,
            text=True,
        )
        self._id = 0

    # -- low-level transport ------------------------------------------------

    def _next_id(self) -> int:
        self._id += 1
        return self._id

    def _send(self, method: str, params: dict[str, Any] | None = None) -> Any:
        """Send a JSON-RPC request and return the result."""
        msg: dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": self._next_id(),
            "method": method,
        }
        if params is not None:
            msg["params"] = params

        payload = json.dumps(msg) + "\n"
        assert self._proc.stdin is not None
        self._proc.stdin.write(payload)
        self._proc.stdin.flush()

        assert self._proc.stdout is not None
        line = self._proc.stdout.readline()
        if not line:
            raise RuntimeError("MCP server closed the connection")
        resp = json.loads(line)
        if "error" in resp:
            raise RuntimeError(f"MCP error: {resp['error']}")
        return resp.get("result")

    def close(self) -> None:
        """Terminate the MCP server process."""
        if self._proc.poll() is None:
            self._proc.terminate()
            self._proc.wait(timeout=5)

    # -- peek-a-bin tool wrappers -------------------------------------------

    def load_pe(self, file_path: str) -> dict[str, Any]:
        """Load and auto-analyze a PE file."""
        return self._send("tools/call", {
            "name": "load_pe",
            "arguments": {"filePath": file_path},
        })

    def detect_anomalies(self, file_id: str) -> dict[str, Any]:
        """Get security anomalies for a loaded PE file."""
        return self._send("tools/call", {
            "name": "detect_anomalies",
            "arguments": {"fileId": file_id},
        })

    def list_functions(
        self,
        file_id: str,
        *,
        filter_name: str | None = None,
        offset: int = 0,
        limit: int = 100,
    ) -> dict[str, Any]:
        """List detected functions with optional filtering."""
        args: dict[str, Any] = {
            "fileId": file_id,
            "offset": offset,
            "limit": limit,
        }
        if filter_name:
            args["filter"] = filter_name
        return self._send("tools/call", {
            "name": "list_functions",
            "arguments": args,
        })

    def decompile_function(self, file_id: str, address: str) -> dict[str, Any]:
        """Decompile a function to C-like pseudocode."""
        return self._send("tools/call", {
            "name": "decompile_function",
            "arguments": {"fileId": file_id, "address": address},
        })

    def disassemble_function(self, file_id: str, address: str) -> dict[str, Any]:
        """Get raw disassembly for a function."""
        return self._send("tools/call", {
            "name": "disassemble_function",
            "arguments": {"fileId": file_id, "address": address},
        })

    def get_xrefs(self, file_id: str, address: str) -> dict[str, Any]:
        """Get cross-references to/from an address."""
        return self._send("tools/call", {
            "name": "get_xrefs",
            "arguments": {"fileId": file_id, "address": address},
        })

    def export_analysis(
        self, file_id: str, output_path: str | None = None
    ) -> dict[str, Any]:
        """Export analysis annotations as JSON."""
        args: dict[str, Any] = {"fileId": file_id}
        if output_path:
            args["outputPath"] = output_path
        return self._send("tools/call", {
            "name": "export_analysis",
            "arguments": args,
        })


# ---------------------------------------------------------------------------
# Exploration workflow
# ---------------------------------------------------------------------------

SUSPICIOUS_FILTERS = [
    "crypt", "encrypt", "decrypt", "http", "url", "socket", "connect",
    "send", "recv", "reg", "registry", "inject", "thread", "process",
    "shell", "exec", "download", "upload", "pipe", "mutex", "service",
]


def _parse_text_content(result: dict[str, Any]) -> str:
    """Extract text from an MCP tool result's content array."""
    if not result:
        return ""
    content = result.get("content", [])
    parts: list[str] = []
    for item in content:
        if isinstance(item, dict) and item.get("type") == "text":
            parts.append(item["text"])
    return "\n".join(parts)


def _parse_json_content(result: dict[str, Any]) -> Any:
    """Extract and parse JSON from an MCP tool result's content array."""
    text = _parse_text_content(result)
    if not text:
        return {}
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return {"raw": text}


def explore(input_path: Path, mcp_command: list[str], mcp_cwd: str | None) -> dict[str, Any]:
    """Run the full exploration workflow and return a structured report."""
    client = McpClient(mcp_command, cwd=mcp_cwd)
    report: dict[str, Any] = {"file": str(input_path), "stages": {}}

    try:
        # Stage 1: Load
        load_result = _parse_json_content(client.load_pe(str(input_path.resolve())))
        report["stages"]["load"] = load_result
        file_id = load_result.get("id", input_path.name)
        entry_point = load_result.get("entryPoint")

        # Stage 2: Anomaly detection
        anomalies = _parse_json_content(client.detect_anomalies(file_id))
        report["stages"]["anomalies"] = anomalies

        # Stage 3: Function survey
        all_functions = _parse_json_content(
            client.list_functions(file_id, limit=200)
        )
        report["stages"]["function_survey"] = {
            "total": all_functions.get("total", 0),
            "sample": all_functions.get("functions", [])[:20],
        }

        # Stage 4: Filtered function search
        interesting: list[dict[str, Any]] = []
        seen_addresses: set[str] = set()
        for term in SUSPICIOUS_FILTERS:
            result = _parse_json_content(
                client.list_functions(file_id, filter_name=term, limit=20)
            )
            for func in result.get("functions", []):
                addr = str(func.get("address", ""))
                if addr and addr not in seen_addresses:
                    seen_addresses.add(addr)
                    interesting.append(func)
        report["stages"]["interesting_functions"] = interesting

        # Stage 5: Decompile entry point + top interesting functions
        decompiled: list[dict[str, Any]] = []
        targets: list[str] = []
        if entry_point:
            targets.append(str(entry_point))
        for func in interesting[:10]:
            addr = str(func.get("address", ""))
            if addr and addr not in targets:
                targets.append(addr)

        for addr in targets[:15]:
            try:
                dec = _parse_json_content(client.decompile_function(file_id, addr))
                decompiled.append({"address": addr, "decompilation": dec})
            except RuntimeError:
                decompiled.append({"address": addr, "error": "decompilation failed"})
        report["stages"]["decompiled"] = decompiled

        # Stage 6: Cross-references for entry point and interesting addresses
        xrefs: list[dict[str, Any]] = []
        for addr in targets[:10]:
            try:
                xr = _parse_json_content(client.get_xrefs(file_id, addr))
                xrefs.append({"address": addr, "xrefs": xr})
            except RuntimeError:
                pass
        report["stages"]["xrefs"] = xrefs

        # Stage 7: Export annotations
        try:
            export = _parse_json_content(client.export_analysis(file_id))
            report["stages"]["export"] = export
        except RuntimeError:
            report["stages"]["export"] = {}

    finally:
        client.close()

    return report


# ---------------------------------------------------------------------------
# Output formatting
# ---------------------------------------------------------------------------

def format_text(report: dict[str, Any]) -> str:
    """Format the report as human-readable text."""
    lines: list[str] = []
    lines.append(f"PE Binary Exploration Report: {report['file']}")
    lines.append("=" * 60)

    load = report["stages"].get("load", {})
    lines.append(f"\nArchitecture: {'x64' if load.get('is64') else 'x86'}")
    lines.append(f"Entry Point:  {load.get('entryPoint', 'unknown')}")
    lines.append(f"Subsystem:    {load.get('subsystem', 'unknown')}")
    lines.append(f"Sections:     {load.get('sectionCount', '?')}")
    lines.append(f"Imports:      {load.get('importCount', '?')}")
    lines.append(f"Functions:    {load.get('functionCount', '?')}")

    anomalies = report["stages"].get("anomalies", {})
    anomaly_list = anomalies.get("anomalies", [])
    lines.append(f"\nAnomalies ({len(anomaly_list)}):")
    for a in anomaly_list:
        lines.append(f"  [{a.get('severity', '?')}] {a.get('title', '?')}")
        if a.get("detail"):
            lines.append(f"         {a['detail']}")

    interesting = report["stages"].get("interesting_functions", [])
    lines.append(f"\nInteresting Functions ({len(interesting)}):")
    for f in interesting[:20]:
        lines.append(f"  {f.get('address', '?'):>12s}  {f.get('name', 'unknown')}")

    decompiled = report["stages"].get("decompiled", [])
    lines.append(f"\nDecompiled Functions ({len(decompiled)}):")
    for d in decompiled:
        dec = d.get("decompilation", {})
        name = dec.get("functionName", d.get("address", "unknown"))
        lines.append(f"\n--- {name} @ {d['address']} ---")
        code = dec.get("code", d.get("error", "no output"))
        lines.append(textwrap.indent(code, "  "))

    return "\n".join(lines)


def format_output(report: dict[str, Any], fmt: str) -> str:
    """Format the exploration report in the requested format."""
    if fmt == "json":
        return json.dumps(report, indent=2)
    if fmt == "text":
        return format_text(report)
    return json.dumps(report, indent=2)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Orchestrate a PE binary exploration workflow via peek-a-bin MCP."
    )
    parser.add_argument(
        "--input", type=Path, required=True,
        help="Path to the PE file to analyze",
    )
    parser.add_argument(
        "--output", type=Path, default=None,
        help="Path to write the report (stdout if omitted)",
    )
    parser.add_argument(
        "--format", default="json", choices=["json", "text"],
        help="Output format (default: json)",
    )
    parser.add_argument(
        "--mcp-command", default="npx tsx src/mcp/index.ts",
        help="Command to launch the peek-a-bin MCP server (default: npx tsx src/mcp/index.ts)",
    )
    parser.add_argument(
        "--mcp-cwd", default=None,
        help="Working directory for the MCP server command (default: current directory)",
    )
    args = parser.parse_args()

    if not args.input.exists():
        print(f"Error: input file not found: {args.input}", file=sys.stderr)
        sys.exit(1)

    mcp_parts = args.mcp_command.split()
    report = explore(args.input, mcp_parts, args.mcp_cwd)
    output = format_output(report, args.format)

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(output, encoding="utf-8")
        print(f"Report written to {args.output}", file=sys.stderr)
    else:
        print(output)


if __name__ == "__main__":
    main()