#!/usr/bin/env python3
"""
capa_runner.py - Wrap capa execution, parse JSON output, and generate readable summaries.

Runs Mandiant's capa tool against PE/ELF/shellcode samples and produces organized
reports grouped by ATT&CK tactic, MBC objective, or capability namespace.
Supports batch analysis of entire directories.

Usage:
    python3 capa_runner.py --input sample.exe
    python3 capa_runner.py --input sample.exe --format markdown --group-by tactic
    python3 capa_runner.py --batch /path/to/samples/ --output results/ --format csv
    python3 capa_runner.py --input sample.exe --rules /path/to/custom/rules/
"""

from __future__ import annotations

import argparse
import csv
import io
import json
import logging
import os
import shutil
import subprocess
import sys
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Capa execution
# ---------------------------------------------------------------------------

def find_capa() -> str:
    """Locate the capa binary."""
    capa_path = shutil.which("capa")
    if capa_path:
        return capa_path
    # Check common locations
    for candidate in ["/usr/local/bin/capa", "/opt/capa/capa", "capa.exe"]:
        if Path(candidate).is_file():
            return candidate
    raise FileNotFoundError(
        "capa not found. Install from https://github.com/mandiant/capa/releases"
    )


def run_capa(
    sample_path: str,
    rules_path: Optional[str] = None,
    capa_binary: Optional[str] = None,
    timeout: int = 300,
) -> dict:
    """Run capa and return parsed JSON output."""
    if not capa_binary:
        capa_binary = find_capa()

    cmd = [capa_binary, "--json", sample_path]
    if rules_path:
        cmd.extend(["--rules", rules_path])

    logger.info(f"Running: {' '.join(cmd)}")

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
    except subprocess.TimeoutExpired:
        logger.error(f"capa timed out after {timeout}s on {sample_path}")
        return {"error": "timeout", "sample": sample_path}
    except FileNotFoundError:
        logger.error(f"capa binary not found: {capa_binary}")
        return {"error": "capa_not_found", "sample": sample_path}

    if result.returncode != 0:
        logger.error(f"capa returned {result.returncode}: {result.stderr.strip()}")
        return {
            "error": result.stderr.strip() or f"exit code {result.returncode}",
            "sample": sample_path,
        }

    try:
        return json.loads(result.stdout)
    except json.JSONDecodeError as e:
        logger.error(f"Failed to parse capa JSON: {e}")
        return {"error": f"JSON parse error: {e}", "sample": sample_path}


# ---------------------------------------------------------------------------
# Result parsing
# ---------------------------------------------------------------------------

def extract_capabilities(capa_result: dict) -> list:
    """Extract capability information from capa JSON output."""
    capabilities = []

    rules = capa_result.get("rules", {})
    for rule_name, rule_data in rules.items():
        meta = rule_data.get("meta", {})

        cap = {
            "name": rule_name,
            "namespace": meta.get("namespace", ""),
            "scope": meta.get("scope", "function"),
            "authors": meta.get("authors", []),
            "description": meta.get("description", ""),
            "references": meta.get("references", []),
            "attack": [],
            "mbc": [],
            "match_count": len(rule_data.get("matches", {})),
        }

        # Extract ATT&CK mappings
        for attack in meta.get("attack", []):
            cap["attack"].append({
                "tactic": attack.get("tactic", ""),
                "technique": attack.get("technique", ""),
                "subtechnique": attack.get("subtechnique", ""),
                "id": attack.get("id", ""),
            })

        # Extract MBC mappings
        for mbc in meta.get("mbc", []):
            cap["mbc"].append({
                "objective": mbc.get("objective", ""),
                "behavior": mbc.get("behavior", ""),
                "method": mbc.get("method", ""),
                "id": mbc.get("id", ""),
            })

        capabilities.append(cap)

    return capabilities


def group_by_tactic(capabilities: list) -> dict:
    """Group capabilities by ATT&CK tactic."""
    groups = defaultdict(list)
    unmapped = []

    for cap in capabilities:
        if cap["attack"]:
            for mapping in cap["attack"]:
                tactic = mapping["tactic"] or "Unknown"
                groups[tactic].append({
                    "capability": cap["name"],
                    "technique_id": mapping["id"],
                    "technique": mapping["technique"],
                    "subtechnique": mapping.get("subtechnique", ""),
                    "namespace": cap["namespace"],
                    "matches": cap["match_count"],
                })
        else:
            unmapped.append(cap)

    if unmapped:
        groups["No ATT&CK Mapping"] = [
            {
                "capability": c["name"],
                "technique_id": "",
                "technique": "",
                "subtechnique": "",
                "namespace": c["namespace"],
                "matches": c["match_count"],
            }
            for c in unmapped
        ]

    return dict(groups)


def group_by_mbc(capabilities: list) -> dict:
    """Group capabilities by MBC objective."""
    groups = defaultdict(list)
    unmapped = []

    for cap in capabilities:
        if cap["mbc"]:
            for mapping in cap["mbc"]:
                objective = mapping["objective"] or "Unknown"
                groups[objective].append({
                    "capability": cap["name"],
                    "behavior_id": mapping["id"],
                    "behavior": mapping["behavior"],
                    "method": mapping.get("method", ""),
                    "namespace": cap["namespace"],
                    "matches": cap["match_count"],
                })
        else:
            unmapped.append(cap)

    if unmapped:
        groups["No MBC Mapping"] = [
            {
                "capability": c["name"],
                "behavior_id": "",
                "behavior": "",
                "method": "",
                "namespace": c["namespace"],
                "matches": c["match_count"],
            }
            for c in unmapped
        ]

    return dict(groups)


def group_by_namespace(capabilities: list) -> dict:
    """Group capabilities by capa namespace."""
    groups = defaultdict(list)
    for cap in capabilities:
        ns = cap["namespace"] or "Uncategorized"
        # Use top-level namespace
        top_ns = ns.split("/")[0] if "/" in ns else ns
        groups[top_ns].append(cap)
    return dict(groups)


def group_by_capability(capabilities: list) -> dict:
    """Flat list grouped by capability name."""
    return {"All Capabilities": capabilities}


GROUPERS = {
    "tactic": group_by_tactic,
    "mbc": group_by_mbc,
    "namespace": group_by_namespace,
    "capability": group_by_capability,
}


# ---------------------------------------------------------------------------
# Output formatters
# ---------------------------------------------------------------------------

def format_summary(capa_result: dict, capabilities: list, grouped: dict, group_by: str) -> str:
    """Format as human-readable summary."""
    lines = []
    meta = capa_result.get("meta", {})
    sample = meta.get("sample", {})

    lines.append("=" * 70)
    lines.append("CAPA ANALYSIS REPORT")
    lines.append("=" * 70)
    lines.append(f"Sample: {sample.get('path', 'N/A')}")
    lines.append(f"MD5:    {sample.get('md5', 'N/A')}")
    lines.append(f"SHA256: {sample.get('sha256', 'N/A')}")
    lines.append(f"Format: {meta.get('analysis', {}).get('format', 'N/A')}")
    lines.append(f"Arch:   {meta.get('analysis', {}).get('arch', 'N/A')}")
    lines.append(f"OS:     {meta.get('analysis', {}).get('os', 'N/A')}")
    lines.append(f"Total capabilities: {len(capabilities)}")
    lines.append(f"Grouped by: {group_by}")
    lines.append("")

    # Count ATT&CK tactics
    all_tactics = set()
    all_techniques = set()
    for cap in capabilities:
        for a in cap["attack"]:
            if a["tactic"]:
                all_tactics.add(a["tactic"])
            if a["id"]:
                all_techniques.add(a["id"])

    lines.append(f"ATT&CK coverage: {len(all_tactics)} tactics, {len(all_techniques)} techniques")
    lines.append("")

    for group_name, items in sorted(grouped.items()):
        lines.append("-" * 50)
        lines.append(f"  {group_name} ({len(items)} capabilities)")
        lines.append("-" * 50)
        for item in items:
            if isinstance(item, dict):
                name = item.get("capability", item.get("name", ""))
                tech = item.get("technique_id", item.get("behavior_id", ""))
                ns = item.get("namespace", "")
                matches = item.get("matches", item.get("match_count", 0))
                detail = f"    [{tech}] " if tech else "    "
                detail += name
                if ns:
                    detail += f"  ({ns})"
                if matches:
                    detail += f"  [{matches} matches]"
                lines.append(detail)

    lines.append("")
    lines.append("=" * 70)
    return "\n".join(lines)


def format_markdown(capa_result: dict, capabilities: list, grouped: dict, group_by: str) -> str:
    """Format as Markdown."""
    lines = []
    meta = capa_result.get("meta", {})
    sample = meta.get("sample", {})

    lines.append("# Capa Analysis Report")
    lines.append("")
    lines.append(f"**Sample**: `{sample.get('path', 'N/A')}`")
    lines.append(f"**MD5**: `{sample.get('md5', 'N/A')}`")
    lines.append(f"**SHA256**: `{sample.get('sha256', 'N/A')}`")
    lines.append(f"**Total capabilities**: {len(capabilities)}")
    lines.append("")

    for group_name, items in sorted(grouped.items()):
        lines.append(f"## {group_name}")
        lines.append("")
        lines.append("| Capability | ID | Namespace | Matches |")
        lines.append("|---|---|---|---|")
        for item in items:
            if isinstance(item, dict):
                name = item.get("capability", item.get("name", ""))
                tech = item.get("technique_id", item.get("behavior_id", ""))
                ns = item.get("namespace", "")
                matches = item.get("matches", item.get("match_count", 0))
                lines.append(f"| {name} | {tech} | {ns} | {matches} |")
        lines.append("")

    return "\n".join(lines)


def format_csv_output(capabilities: list) -> str:
    """Format as CSV."""
    output = io.StringIO()
    writer = csv.writer(output)
    writer.writerow([
        "capability", "namespace", "scope", "match_count",
        "attack_tactic", "attack_technique", "attack_id",
        "mbc_objective", "mbc_behavior", "mbc_id",
    ])
    for cap in capabilities:
        attack_rows = cap["attack"] if cap["attack"] else [{}]
        mbc_rows = cap["mbc"] if cap["mbc"] else [{}]

        for atk in attack_rows:
            for mbc in mbc_rows:
                writer.writerow([
                    cap["name"],
                    cap["namespace"],
                    cap["scope"],
                    cap["match_count"],
                    atk.get("tactic", ""),
                    atk.get("technique", ""),
                    atk.get("id", ""),
                    mbc.get("objective", ""),
                    mbc.get("behavior", ""),
                    mbc.get("id", ""),
                ])
    return output.getvalue()


FORMATTERS = {
    "summary": format_summary,
    "markdown": format_markdown,
    "json": None,  # Handled separately
    "csv": None,   # Handled separately
}


# ---------------------------------------------------------------------------
# Batch processing
# ---------------------------------------------------------------------------

SAMPLE_EXTENSIONS = {".exe", ".dll", ".sys", ".bin", ".elf", ".so", ".dylib", ".scr", ".cpl"}


def process_batch(
    batch_dir: str,
    output_dir: str,
    rules_path: Optional[str],
    group_by: str,
    fmt: str,
    capa_binary: Optional[str],
) -> None:
    """Process all samples in a directory."""
    batch_path = Path(batch_dir)
    out_path = Path(output_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    samples = []
    for f in batch_path.iterdir():
        if f.is_file() and (f.suffix.lower() in SAMPLE_EXTENSIONS or not f.suffix):
            samples.append(f)

    if not samples:
        logger.error(f"No samples found in {batch_dir}")
        return

    logger.info(f"Found {len(samples)} samples in {batch_dir}")

    results_summary = []
    for i, sample in enumerate(sorted(samples), 1):
        logger.info(f"[{i}/{len(samples)}] Analyzing {sample.name}...")
        capa_result = run_capa(str(sample), rules_path, capa_binary)

        if "error" in capa_result:
            logger.warning(f"  Error: {capa_result['error']}")
            results_summary.append({
                "sample": sample.name,
                "status": "error",
                "error": capa_result["error"],
                "capabilities": 0,
            })
            continue

        capabilities = extract_capabilities(capa_result)
        results_summary.append({
            "sample": sample.name,
            "status": "success",
            "capabilities": len(capabilities),
            "attack_techniques": len({a["id"] for c in capabilities for a in c["attack"] if a.get("id")}),
        })

        # Write per-sample report
        sample_out = out_path / f"{sample.stem}_capa.{_extension(fmt)}"
        output = _format_output(capa_result, capabilities, group_by, fmt)
        sample_out.write_text(output)
        logger.info(f"  {len(capabilities)} capabilities -> {sample_out}")

    # Write batch summary
    summary_path = out_path / f"batch_summary.json"
    summary_path.write_text(json.dumps({
        "batch_dir": str(batch_dir),
        "timestamp": datetime.now().isoformat(),
        "total_samples": len(samples),
        "successful": sum(1 for r in results_summary if r["status"] == "success"),
        "failed": sum(1 for r in results_summary if r["status"] == "error"),
        "results": results_summary,
    }, indent=2))
    logger.info(f"Batch summary: {summary_path}")


def _extension(fmt: str) -> str:
    return {"summary": "txt", "markdown": "md", "json": "json", "csv": "csv"}.get(fmt, "txt")


def _format_output(capa_result: dict, capabilities: list, group_by: str, fmt: str) -> str:
    if fmt == "json":
        return json.dumps(capa_result, indent=2)
    if fmt == "csv":
        return format_csv_output(capabilities)

    grouper = GROUPERS.get(group_by, group_by_tactic)
    grouped = grouper(capabilities)

    if fmt == "markdown":
        return format_markdown(capa_result, capabilities, grouped, group_by)
    return format_summary(capa_result, capabilities, grouped, group_by)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run capa analysis and generate organized capability reports.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --input sample.exe
  %(prog)s --input sample.exe --format markdown --group-by tactic
  %(prog)s --input sample.exe --format json --output report.json
  %(prog)s --batch /samples/ --output results/ --format csv
  %(prog)s --input sample.exe --rules /path/to/custom/rules/
        """,
    )

    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument("--input", "-i", help="Single sample file to analyze")
    input_group.add_argument("--batch", "-b", help="Directory of samples for batch analysis")

    parser.add_argument("--output", "-o", help="Output file or directory (for batch)")
    parser.add_argument(
        "--format", "-f",
        choices=["summary", "json", "csv", "markdown"],
        default="summary",
        help="Output format (default: summary)",
    )
    parser.add_argument(
        "--group-by", "-g",
        choices=["tactic", "mbc", "namespace", "capability"],
        default="tactic",
        help="Group results by category (default: tactic)",
    )
    parser.add_argument("--rules", "-r", help="Path to custom capa rules directory")
    parser.add_argument("--capa-path", help="Path to capa binary (auto-detected by default)")
    parser.add_argument("--timeout", type=int, default=300, help="Capa execution timeout in seconds (default: 300)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    return parser.parse_args()


def main() -> None:
    args = parse_arguments()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s: %(message)s",
    )

    # Batch mode
    if args.batch:
        output_dir = args.output or "capa_batch_results"
        process_batch(args.batch, output_dir, args.rules, args.group_by, args.format, args.capa_path)
        return

    # Single file mode
    sample_path = args.input
    if not Path(sample_path).is_file():
        logger.error(f"Sample not found: {sample_path}")
        sys.exit(1)

    capa_result = run_capa(sample_path, args.rules, args.capa_path, args.timeout)

    if "error" in capa_result:
        logger.error(f"capa analysis failed: {capa_result['error']}")
        sys.exit(1)

    capabilities = extract_capabilities(capa_result)
    output = _format_output(capa_result, capabilities, args.group_by, args.format)

    if args.output:
        Path(args.output).write_text(output)
        logger.info(f"Report written to {args.output}")
    else:
        print(output)


if __name__ == "__main__":
    main()