"""
Multi-Agent Reverse Engineering System
====================================
A team of specialized agents, each with its own strength, coordinated by an
orchestrator. Instead of one monolithic agent, tasks are routed to the agent
best suited for them, and their findings are fused into a single report.

Agents:
  - BinaryAgent        native code (ELF/PE/Mach-O): disassembly, decompilation, CFG
  - WebAgent           websites / JS bundles / SPAs: endpoints, secrets, source maps
  - MobileAgent        APK/IPA: manifests, permissions, components, native libs
  - FirmwareAgent      raw blobs: entropy, partitions, strings, arch guessing
  - MalwareAgent       IOCs, anti-analysis, persistence, C2, crypto misuse
  - ReconstructorAgent rebuilds a project scaffold from the fused analysis
  - ReporterAgent      synthesizes the final report

Usage:
    python multi_agent.py analyze ./target.bin
    python multi_agent.py analyze https://example.com --mode web
    python multi_agent.py analyze app.apk --mode apk
"""

from __future__ import annotations
import os, json, hashlib, shutil, tempfile, zipfile, subprocess
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import List, Dict, Any, Optional

import agent

@dataclass
class AgentResult:
    agent: str
    strength: str
    findings: Dict[str, Any] = field(default_factory=dict)
    errors: List[str] = field(default_factory=list)
    confidence: float = 0.0

@dataclass
class FusedReport:
    target: str
    mode: str
    agents_run: List[str] = field(default_factory=list)
    results: List[AgentResult] = field(default_factory=list)
    summary: str = ""
    iocs: List[str] = field(default_factory=list)
    risks: List[str] = field(default_factory=list)
    rebuild_path: Optional[str] = None
    report: str = ""

class BaseAgent:
    name: str = "base"
    strength: str = "general"

    def can_handle(self, ctx: agent.AnalysisContext) -> bool:
        return True

    def run(self, ctx: agent.AnalysisContext) -> AgentResult:
        raise NotImplementedError

    def _result(self, findings, errors=None, confidence=0.5):
        return AgentResult(agent=self.name, strength=self.strength,
                           findings=findings, errors=errors or [], confidence=confidence)

class BinaryAgent(BaseAgent):
    name = "BinaryAgent"
    strength = "native code: disassembly, decompilation, CFG, symbols"

    def can_handle(self, ctx):
        return ctx.mode == "binary"

    def run(self, ctx):
        errs = []
        ctx = agent.FileIdentificationStage().run(ctx)
        ctx = agent.StringExtractionStage(min_len=4).run(ctx)
        ctx = agent.DisassemblyStage().run(ctx)
        findings = {
            "format": ctx.format, "arch": ctx.arch, "bits": ctx.bits,
            "sections": len(ctx.sections), "imports": len(ctx.imports),
            "exports": len(ctx.exports), "strings": len(ctx.strings),
            "disassembled_sections": len(ctx.disassembly),
        }
        if shutil.which("analyzeHeadless") or shutil.which("r2"):
            agent.DecompilationStage().run(ctx)
            findings["decompiled"] = len(ctx.decompiled)
        return self._result(findings, errs, 0.9)

class WebAgent(BaseAgent):
    name = "WebAgent"
    strength = "web apps: endpoints, JS, secrets, source maps, forms"

    def can_handle(self, ctx):
        return ctx.mode == "web"

    def run(self, ctx):
        ctx = agent.WebFetchStage().run(ctx)
        ctx = agent.WebAnalyzeStage().run(ctx)
        findings = {
            "endpoints": len(ctx.endpoints),
            "secrets": len(ctx.secrets),
            "strings": len(ctx.strings),
            "forms": len(ctx.metadata.get("forms", [])),
            "scripts": len(ctx.metadata.get("scripts", [])),
        }
        return self._result(findings, ctx.errors, 0.85)

class MobileAgent(BaseAgent):
    name = "MobileAgent"
    strength = "mobile apps: manifests, permissions, components, native libs"

    def can_handle(self, ctx):
        return ctx.mode in ("apk", "ipa")

    def run(self, ctx):
        ctx = agent.ApkUnzipStage().run(ctx)
        ctx = agent.ApkAnalyzeStage().run(ctx)
        findings = {
            "assets": len(ctx.assets),
            "permissions": len(ctx.permissions),
            "components": len(ctx.components),
            "secrets": len(ctx.secrets),
            "endpoints": len(ctx.endpoints),
        }
        return self._result(findings, ctx.errors, 0.85)

class FirmwareAgent(BaseAgent):
    name = "FirmwareAgent"
    strength = "raw firmware: entropy, partition hints, arch guessing, strings"

    def can_handle(self, ctx):
        return ctx.mode == "binary" and ctx.format == "raw"

    def run(self, ctx):
        data = ctx.target.read_bytes()
        import math
        from collections import Counter
        def entropy(b):
            if not b: return 0.0
            c = Counter(b); n = len(b)
            return -sum((v/n)*math.log2(v/n) for v in c.values())
        blocks = [data[i:i+4096] for i in range(0, len(data), 4096)]
        ent = [round(entropy(b), 3) for b in blocks]
        high = [i for i, e in enumerate(ent) if e > 7.5]
        ctx = agent.StringExtractionStage(min_len=6).run(ctx)
        findings = {
            "size": len(data),
            "blocks": len(blocks),
            "avg_entropy": round(sum(ent)/max(1,len(ent)), 3),
            "high_entropy_blocks": len(high),
            "strings": len(ctx.strings),
            "urls": len([s for s in ctx.strings if "url" in s.get("tags", [])]),
        }
        ctx.metadata["entropy"] = ent
        return self._result(findings, ctx.errors, 0.7)

class MalwareAgent(BaseAgent):
    name = "MalwareAgent"
    strength = "threat intel: IOCs, C2, persistence, anti-analysis, crypto misuse"

    def can_handle(self, ctx):
        return True

    SUSPICIOUS_IMPORTS = {"system", "popen", "execve", "VirtualAlloc", "WriteProcessMemory",
                          "CreateRemoteThread", "RegSetValueEx", "connect", "socket", "send"}
    PERSIST_PATHS = ["/etc/cron", "HKCU\\Software\\Microsoft\\Windows\\CurrentVersion\\Run",
                     "LaunchAgents", "systemd", "init.d"]

    def run(self, ctx):
        iocs = []
        risks = []
        for imp in ctx.imports:
            if imp.get("name") in self.SUSPICIOUS_IMPORTS:
                iocs.append(f"suspicious import: {imp['name']}")
        for s in ctx.strings:
            v = s.get("value", "")
            if "http://" in v or "https://" in v:
                iocs.append(f"url: {v}")
            for p in self.PERSIST_PATHS:
                if p in v:
                    risks.append(f"persistence path: {v}")
            if any(k in v.lower() for k in ("password", "secret", "token", "api_key")):
                risks.append(f"credential-like string: {v[:60]}")
        for s in ctx.strings:
            v = s.get("value", "").lower()
            if any(k in v for k in ("isdebuggerpresent", "ptrace", "sandbox", "vmware", "virtualbox")):
                risks.append(f"anti-analysis: {s['value']}")
        findings = {"iocs": len(iocs), "risks": len(risks),
                    "sample_iocs": iocs[:20], "sample_risks": risks[:20]}
        ctx.metadata["malware"] = findings
        return self._result(findings, ctx.errors, 0.75)

class ReconstructorAgent(BaseAgent):
    name = "ReconstructorAgent"
    strength = "rebuilds a project scaffold from fused analysis"

    def can_handle(self, ctx):
        return True

    def run(self, ctx):
        out = Path(tempfile.mkdtemp()) / "rebuilt"
        out.mkdir()
        if ctx.mode == "binary":
            (out / "src").mkdir()
            (out / "src" / "main.c").write_text(self._c_scaffold(ctx))
            (out / "Makefile").write_text("all:\n\tgcc src/main.c -o app\n")
        elif ctx.mode == "web":
            (out / "index.html").write_text(ctx.metadata.get("html", "<!-- captured html -->"))
            (out / "endpoints.txt").write_text("\n".join(ctx.endpoints))
            (out / "secrets.txt").write_text("\n".join(f"{s.get('type')}: {s.get('value')}" for s in ctx.secrets))
        else:
            extracted = ctx.metadata.get("extracted")
            if extracted and Path(extracted).exists():
                shutil.copytree(extracted, out / "extracted", dirs_exist_ok=True)
            (out / "permissions.txt").write_text("\n".join(ctx.permissions))
            (out / "components.json").write_text(json.dumps(ctx.components, indent=2))
            (out / "secrets.txt").write_text("\n".join(f"{s.get('type')}: {s.get('value')}" for s in ctx.secrets))
        (out / "README.md").write_text(f"# Rebuilt from {ctx.target}\n\nMode: {ctx.mode}\n")
        zip_path = str(out.parent / "rebuilt.zip")
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
            for f in out.rglob("*"):
                if f.is_file():
                    z.write(f, f.relative_to(out.parent))
        return self._result({"rebuild_path": zip_path}, ctx.errors, 0.6)

    def _c_scaffold(self, ctx):
        lines = ["#include <stdio.h>", "#include <stdlib.h>", "", "/* Reconstructed from reverse engineering. Fill in bodies. */", ""]
        for f in getattr(ctx, "functions", []) or []:
            lines.append(f"/* {f.name}: {f.summary} */")
            lines.append(f"void {f.name}(void) {{ /* TODO */ }}")
            lines.append("")
        if not getattr(ctx, "functions", []):
            lines.append("int main(int argc, char** argv) {")
            lines.append("    printf(\"rebuilt\\n\");")
            lines.append("    return 0;")
            lines.append("}")
        return "\n".join(lines)

class ReporterAgent(BaseAgent):
    name = "ReporterAgent"
    strength = "synthesis: final report, risk summary, IOC list"

    def can_handle(self, ctx):
        return True

    def run(self, ctx):
        iocs = ctx.metadata.get("malware", {}).get("sample_iocs", [])
        risks = ctx.metadata.get("malware", {}).get("sample_risks", [])
        lines = ["# Multi-Agent Reverse Engineering Report", "", f"- Target: `{ctx.target}`", f"- Mode: {ctx.mode}", ""]
        lines.append("## Agent findings")
        lines.append(f"- strings: {len(ctx.strings)}")
        lines.append(f"- endpoints: {len(ctx.endpoints)}")
        lines.append(f"- secrets: {len(ctx.secrets)}")
        lines.append(f"- permissions: {len(ctx.permissions)}")
        lines.append(f"- components: {len(ctx.components)}")
        if iocs:
            lines.append("")
            lines.append("## IOCs")
            lines += [f"- {i}" for i in iocs]
        if risks:
            lines.append("")
            lines.append("## Risks")
            lines += [f"- {r}" for r in risks]
        if ctx.strings:
            lines.append("")
            lines.append("## Strings (sample)")
            lines += [f"- {s.get('value','')[:120]}" for s in ctx.strings[:60]]
        ctx.report = "\n".join(lines)
        return self._result({"report_len": len(ctx.report), "iocs": len(iocs), "risks": len(risks)}, ctx.errors, 0.9)

AGENTS = [BinaryAgent(), WebAgent(), MobileAgent(), FirmwareAgent(), MalwareAgent(),
          ReconstructorAgent(), ReporterAgent()]

class Orchestrator:
    def __init__(self, mode="binary", model="gpt-4o-mini", rebuild=True):
        self.mode = mode
        self.model = model
        self.rebuild = rebuild

    def run(self, target) -> FusedReport:
        ctx = agent.AnalysisContext(target=Path(target) if not str(target).startswith("http") else Path(target),
                                    mode=self.mode)
        if self.mode == "web":
            ctx.target = Path(target)
        fr = FusedReport(target=str(target), mode=self.mode)
        for ag in AGENTS:
            if not ag.can_handle(ctx):
                continue
            if ag.name in ("ReconstructorAgent", "ReporterAgent"):
                continue
            try:
                res = ag.run(ctx)
                fr.results.append(res)
                fr.agents_run.append(ag.name)
            except Exception as e:
                fr.results.append(AgentResult(agent=ag.name, strength=ag.strength, errors=[str(e)]))
        if "MalwareAgent" not in fr.agents_run:
            try:
                res = MalwareAgent().run(ctx)
                fr.results.append(res); fr.agents_run.append("MalwareAgent")
            except Exception as e:
                fr.results.append(AgentResult(agent="MalwareAgent", strength="", errors=[str(e)]))
        agent.AIAugmentationStage(model=self.model).run(ctx)
        if self.rebuild:
            try:
                res = ReconstructorAgent().run(ctx)
                fr.results.append(res); fr.agents_run.append("ReconstructorAgent")
                fr.rebuild_path = res.findings.get("rebuild_path")
            except Exception as e:
                fr.results.append(AgentResult(agent="ReconstructorAgent", strength="", errors=[str(e)]))
        try:
            ReporterAgent().run(ctx)
            fr.agents_run.append("ReporterAgent")
        except Exception as e:
            fr.results.append(AgentResult(agent="ReporterAgent", strength="", errors=[str(e)]))
        fr.summary = f"Ran {len(fr.agents_run)} agents on {target}"
        fr.iocs = ctx.metadata.get("malware", {}).get("sample_iocs", [])
        fr.risks = ctx.metadata.get("malware", {}).get("sample_risks", [])
        fr.report = ctx.report
        return fr

    def to_dict(self, fr):
        return {
            "target": fr.target, "mode": fr.mode,
            "agents_run": fr.agents_run,
            "results": [asdict(r) for r in fr.results],
            "iocs": fr.iocs, "risks": fr.risks,
            "rebuild_path": fr.rebuild_path,
        }

def main():
    import argparse
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    a = sub.add_parser("analyze")
    a.add_argument("target")
    a.add_argument("--mode", default="binary", choices=["binary", "web", "apk", "ipa"])
    a.add_argument("--out", default="report.md")
    a.add_argument("--json", default=None)
    a.add_argument("--no-rebuild", action="store_true")
    args = ap.parse_args()
    if args.cmd == "analyze":
        orch = Orchestrator(mode=args.mode, rebuild=not args.no_rebuild)
        fr = orch.run(args.target)
        Path(args.out).write_text(fr.report)
        if args.json:
            Path(args.json).write_text(json.dumps(orch.to_dict(fr), indent=2, default=str))
        print(f"agents: {', '.join(fr.agents_run)}")
        print(f"report -> {args.out}")
        if fr.rebuild_path:
            print(f"rebuild -> {fr.rebuild_path}")

if __name__ == "__main__":
    from pathlib import Path
    main()