agentre_triage.py

#!/usr/bin/env python3
"""User-facing AgentRE PE/ELF triage CLI.

This is intentionally not an eval harness and does not require vLLM. It loads a
local Qwen/Qwen3.5-style model with an optional LoRA adapter through
Transformers, exposes static reverse-engineering tools, and lets the model
inspect one file or every PE/ELF file in a directory.

Static analysis only: samples are copied to neutral staging paths, never
executed, and tool path arguments are ignored.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parent
STARTER = ROOT.parents[1] / "agentre_rl_starter"

DEFAULT_BASE_MODEL = os.environ.get("XREF9B_BASE_MODEL") or os.environ.get("AGENTRE_BASE_MODEL") or "Qwen/Qwen3.5-9B"
DEFAULT_ADAPTER = os.environ.get("XREF9B_ADAPTER") or os.environ.get("AGENTRE_ADAPTER", "")
DEFAULT_SPEC = ROOT / "reverse_engineering_spec_unified.md"
DEFAULT_GHIDRA_SCRIPT_DIR = ROOT / "tools" / "ghidra_scripts"
DEFAULT_LLAMA_CLI = os.environ.get("AGENTRE_LLAMA_CLI", "llama-completion")

FINAL_TOOLS = {"final_answer", "submit_answer", "submit"}

TOOL_LIMITS = {
    "file": 3000,
    "strings": 9000,
    "readelf": 14000,
    "objdump": 18000,
    "nm": 8000,
    "hexdump": 8000,
    "xxd": 8000,
    "entropy": 2000,
    "disasm_func": 18000,
    "pe_headers": 14000,
    "pe_sections": 9000,
    "pe_imports": 14000,
    "pe_exports": 9000,
    "pe_disasm": 18000,
    "pe_symbols": 9000,
    "ghidra_summary": 24000,
}

TOOL_ALIASES = {
    "file": "file",
    "run_file": "file",
    "strings": "strings",
    "run_strings": "strings",
    "readelf": "readelf",
    "run_readelf": "readelf",
    "objdump": "objdump",
    "run_objdump": "objdump",
    "nm": "nm",
    "run_nm": "nm",
    "hexdump": "hexdump",
    "run_hexdump": "hexdump",
    "xxd": "xxd",
    "run_xxd": "xxd",
    "entropy": "entropy",
    "run_entropy": "entropy",
    "disasm_func": "disasm_func",
    "run_disasm_func": "disasm_func",
    "pe_headers": "pe_headers",
    "run_pe_headers": "pe_headers",
    "pe_sections": "pe_sections",
    "run_pe_sections": "pe_sections",
    "pe_imports": "pe_imports",
    "run_pe_imports": "pe_imports",
    "pe_exports": "pe_exports",
    "run_pe_exports": "pe_exports",
    "pe_disasm": "pe_disasm",
    "run_pe_disasm": "pe_disasm",
    "pe_symbols": "pe_symbols",
    "run_pe_symbols": "pe_symbols",
    "ghidra_summary": "ghidra_summary",
    "run_ghidra": "ghidra_summary",
}

XML_TOOL_RE = re.compile(
    r"<tool_call>\s*<function=([^>\n]+)>\s*(.*?)\s*</function>\s*</tool_call>",
    re.DOTALL,
)
XML_PARAM_RE = re.compile(
    r"<parameter=([^>\n]+)>\s*(.*?)\s*</parameter>",
    re.DOTALL,
)
JSON_TOOL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
THINK_RE = re.compile(r"<think>\s*(.*?)\s*</think>\s*", re.DOTALL)


def _now_stamp() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")


def _json_from(text: str, pos: int = 0) -> Any | None:
    decoder = json.JSONDecoder()
    i = text.find("{", pos)
    while i >= 0:
        try:
            obj, _ = decoder.raw_decode(text[i:])
            return obj
        except Exception:
            i = text.find("{", i + 1)
    return None


def _parse_json_maybe(value: Any) -> Any:
    if not isinstance(value, str):
        return value
    text = value.strip()
    if text.startswith("```"):
        lines = text.splitlines()
        if lines and lines[0].lstrip().startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip().startswith("```"):
            lines = lines[:-1]
        text = "\n".join(lines).strip()
    try:
        return json.loads(text)
    except Exception:
        obj = _json_from(text)
        return obj if obj is not None else value


def _parse_args(raw: Any) -> dict[str, Any]:
    if isinstance(raw, dict):
        return raw
    if raw in (None, ""):
        return {}
    parsed = _parse_json_maybe(raw)
    return parsed if isinstance(parsed, dict) else {}


def _tool_call(name: str, args: Any, idx: int) -> dict[str, Any]:
    if not isinstance(args, dict):
        args = {}
    return {
        "id": f"call_{idx}",
        "type": "function",
        "function": {"name": name, "arguments": args},
    }


def parse_tool_calls(text: str | None) -> list[dict[str, Any]]:
    content = text or ""
    calls: list[dict[str, Any]] = []

    for idx, match in enumerate(XML_TOOL_RE.finditer(content)):
        name = match.group(1).strip()
        body = match.group(2)
        args: dict[str, Any] = {}
        for param in XML_PARAM_RE.finditer(body):
            key = param.group(1).strip()
            value = param.group(2).strip()
            args[key] = _parse_json_maybe(value)
        calls.append(_tool_call(name, args, idx))
    if calls:
        return calls

    for idx, match in enumerate(JSON_TOOL_RE.finditer(content)):
        try:
            obj = json.loads(match.group(1))
        except Exception:
            continue
        name = obj.get("name") or obj.get("function", {}).get("name")
        args = obj.get("arguments") or obj.get("function", {}).get("arguments") or {}
        if name:
            calls.append(_tool_call(str(name), _parse_args(args), idx))
    if calls:
        return calls

    lower = content.lower()
    for marker in ("final_answer", "submit_answer", "submit"):
        pos = lower.find(marker)
        if pos >= 0:
            obj = _json_from(content, pos)
            if isinstance(obj, dict):
                if obj.get("name") in FINAL_TOOLS:
                    return [_tool_call(str(obj["name"]), obj.get("arguments", {}), 0)]
                return [_tool_call(marker, obj, 0)]

    obj = _json_from(content)
    if isinstance(obj, dict) and (
        "classification" in obj
        or "recommended_label" in obj
        or "malware" in obj
        or "hackware" in obj
    ):
        return [_tool_call("final_answer", {"answer_json": obj}, 0)]
    return []


def split_thinking(text: str) -> tuple[str | None, str]:
    match = THINK_RE.search(text)
    if match:
        reasoning = match.group(1).strip()
        visible = text[: match.start()] + text[match.end() :]
        return reasoning, visible.strip()

    # llama.cpp completion receives the opening <think> in the prompt, so the
    # generated text may contain only the closing tag. Treat the prefix as
    # hidden reasoning in that case.
    end_tag = "</think>"
    pos = text.find(end_tag)
    if pos >= 0:
        reasoning = text[:pos].strip()
        visible = text[pos + len(end_tag) :].strip()
        return reasoning or None, visible
    return None, text


def strip_tool_xml(text: str) -> str:
    text = XML_TOOL_RE.sub("", text)
    text = JSON_TOOL_RE.sub("", text)
    return text.strip()


def final_answer_from_args(args: dict[str, Any]) -> dict[str, Any] | None:
    answer = args.get("answer_json", args)
    answer = _parse_json_maybe(answer)
    return answer if isinstance(answer, dict) else None


def normalize_prediction(answer: Any) -> str:
    if not isinstance(answer, dict):
        return "unknown"
    for key in ("classification", "recommended_label", "label", "verdict"):
        value = answer.get(key)
        if isinstance(value, str):
            return normalize_prediction_string(value)
    if bool(answer.get("hackware")):
        return "hackware"
    if isinstance(answer.get("malware"), bool):
        return "malicious" if answer["malware"] else "benign"
    return normalize_prediction_string(json.dumps(answer, sort_keys=True))


def normalize_prediction_string(value: str) -> str:
    text = value.lower()
    if "hackware" in text or "dual-use" in text or "offensive tool" in text:
        return "hackware"
    if "benign" in text or "not malware" in text or "non-malicious" in text:
        return "benign"
    if "malicious" in text or "malware" in text or "c2" in text or "reverse shell" in text:
        return "malicious"
    return "unknown"


def openai_tool(name: str, description: str, props: dict[str, Any] | None = None) -> dict[str, Any]:
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": description,
            "parameters": {
                "type": "object",
                "properties": props or {},
                "required": [],
            },
        },
    }


def tool_definitions() -> list[dict[str, Any]]:
    path_prop = {"path": {"type": "string", "description": "Optional. Ignored; routed to current sample."}}
    str_props = dict(path_prop)
    str_props["min_length"] = {"type": "integer", "description": "Minimum printable string length."}
    dump_props = dict(path_prop)
    dump_props["length"] = {"type": "integer", "description": "Maximum bytes to dump."}
    dump_props["offset"] = {"type": "integer", "description": "Start offset in bytes."}
    disasm_props = {"function": {"type": "string", "description": "Function symbol name to disassemble."}}
    ghidra_props = {
        "timeout": {"type": "integer", "description": "Optional max seconds for Ghidra analysis."}
    }
    final_props = {"answer_json": {"type": "object", "description": "Final JSON verdict."}}
    T = openai_tool
    return [
        T("run_file", "Identify file type, format, architecture, linkage, and stripped status.", path_prop),
        T("file", "Alias for run_file.", path_prop),
        T("run_strings", "Extract printable strings from the sample.", str_props),
        T("strings", "Alias for run_strings.", str_props),
        T("run_readelf", "Inspect ELF headers, sections, symbols, program headers, and dynamic imports.", path_prop),
        T("readelf", "Alias for run_readelf.", path_prop),
        T("run_objdump", "Disassemble or dump ELF sections with objdump.", path_prop),
        T("objdump", "Alias for run_objdump.", path_prop),
        T("nm", "List symbols when present.", path_prop),
        T("run_pe_headers", "PE file and optional headers.", path_prop),
        T("pe_headers", "Alias for run_pe_headers.", path_prop),
        T("run_pe_sections", "PE section table and entropy.", path_prop),
        T("pe_sections", "Alias for run_pe_sections.", path_prop),
        T("run_pe_imports", "PE import tables and WinAPI symbols.", path_prop),
        T("pe_imports", "Alias for run_pe_imports.", path_prop),
        T("run_pe_exports", "PE export table and data directories.", path_prop),
        T("pe_exports", "Alias for run_pe_exports.", path_prop),
        T("run_pe_disasm", "Disassemble a PE sample.", path_prop),
        T("pe_disasm", "Alias for run_pe_disasm.", path_prop),
        T("run_pe_symbols", "List PE symbols when present.", path_prop),
        T("pe_symbols", "Alias for run_pe_symbols.", path_prop),
        T("run_disasm_func", "Disassemble one named function.", disasm_props),
        T("disasm_func", "Alias for run_disasm_func.", disasm_props),
        T("hexdump", "Hex dump bytes from the sample.", dump_props),
        T("xxd", "Alias hex dump using xxd.", dump_props),
        T("entropy", "Compute whole-file Shannon entropy.", path_prop),
        T("ghidra_summary", "Run Ghidra headless static analysis and summarize output.", ghidra_props),
        T("final_answer", "Submit exactly one final JSON verdict.", final_props),
        T("submit_answer", "Alias for final_answer.", final_props),
    ]


def render_qwen_xml_tool_call(name: str, args: dict[str, Any]) -> str:
    parts = ["<tool_call>\n", f"<function={name}>\n"]
    for key, value in args.items():
        if isinstance(value, (dict, list)):
            value_text = json.dumps(value, ensure_ascii=True)
        else:
            value_text = str(value)
        parts.extend([f"<parameter={key}>\n", value_text, "\n</parameter>\n"])
    parts.append("</function>\n</tool_call>")
    return "".join(parts)


def render_qwen_xml_chat(
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
    thinking: bool,
    add_generation_prompt: bool = True,
) -> str:
    """Render a Qwen3.5 XML-tool prompt for llama.cpp GGUF inference."""
    chunks: list[str] = []
    system_content = ""
    start_index = 0
    if messages and messages[0].get("role") == "system":
        system_content = str(messages[0].get("content") or "").strip()
        start_index = 1

    if tools:
        chunks.append("<|im_start|>system\n")
        chunks.append("# Tools\n\nYou have access to the following functions:\n\n<tools>")
        for tool in tools:
            chunks.append("\n")
            chunks.append(json.dumps(tool, ensure_ascii=True))
        chunks.append("\n</tools>")
        chunks.append(
            "\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n"
            "<tool_call>\n<function=example_function_name>\n"
            "<parameter=example_parameter_1>\nvalue_1\n</parameter>\n"
            "<parameter=example_parameter_2>\nThis is the value for the second parameter\n"
            "that can span\nmultiple lines\n</parameter>\n"
            "</function>\n</tool_call>\n\n"
            "<IMPORTANT>\n"
            "Reminder:\n"
            "- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n"
            "- Required parameters MUST be specified\n"
            "- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n"
            "- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n"
            "</IMPORTANT>"
        )
        if system_content:
            chunks.append("\n\n")
            chunks.append(system_content)
        chunks.append("<|im_end|>\n")
    elif system_content:
        chunks.append("<|im_start|>system\n")
        chunks.append(system_content)
        chunks.append("<|im_end|>\n")

    i = start_index
    while i < len(messages):
        msg = messages[i]
        role = msg.get("role")
        content = str(msg.get("content") or "").strip()
        if role == "user":
            chunks.append("<|im_start|>user\n")
            chunks.append(content)
            chunks.append("<|im_end|>\n")
        elif role == "assistant":
            chunks.append("<|im_start|>assistant\n")
            reasoning = str(msg.get("reasoning_content") or "").strip()
            if reasoning:
                chunks.append("<think>\n")
                chunks.append(reasoning)
                chunks.append("\n</think>\n\n")
            elif msg.get("reasoning_content") is not None:
                chunks.append("<think>\n\n</think>\n\n")
            chunks.append(content)
            tool_calls = msg.get("tool_calls") or []
            if tool_calls:
                if content:
                    chunks.append("\n\n")
                for tc in tool_calls:
                    fn = tc.get("function", {})
                    name = str(fn.get("name") or "")
                    args = _parse_args(fn.get("arguments"))
                    chunks.append(render_qwen_xml_tool_call(name, args))
                    chunks.append("\n")
            chunks.append("<|im_end|>\n")
        elif role == "tool":
            chunks.append("<|im_start|>user")
            while i < len(messages) and messages[i].get("role") == "tool":
                tool_msg = messages[i]
                chunks.append("\n<tool_response>\n")
                chunks.append(str(tool_msg.get("content") or ""))
                chunks.append("\n</tool_response>")
                i += 1
            chunks.append("<|im_end|>\n")
            continue
        i += 1

    if add_generation_prompt:
        chunks.append("<|im_start|>assistant\n")
        if thinking:
            chunks.append("<think>\n")
        else:
            chunks.append("<think>\n\n</think>\n\n")
    return "".join(chunks)


def _parse_int(value: Any, default: int) -> int:
    try:
        if isinstance(value, str):
            return int(value.strip(), 0)
        return int(value)
    except Exception:
        return default


def _dump_bounds(length: Any, offset: Any) -> tuple[int, int]:
    length_i = _parse_int(length, 1024)
    offset_i = _parse_int(offset, 0)
    return max(64, min(length_i, 8192)), max(0, offset_i)


def detect_format(path: Path) -> str:
    try:
        magic = path.read_bytes()[:4]
    except Exception:
        return "unknown"
    if magic == b"\x7fELF":
        return "ELF"
    if magic[:2] == b"MZ":
        return "PE"
    return "unknown"


def collect_targets(path: Path, include_unknown: bool, max_files: int) -> list[Path]:
    if path.is_file():
        return [path]
    if not path.is_dir():
        raise FileNotFoundError(path)
    out: list[Path] = []
    for p in sorted(path.rglob("*")):
        if not p.is_file() or p.is_symlink():
            continue
        fmt = detect_format(p)
        if fmt in {"PE", "ELF"} or include_unknown:
            out.append(p)
        if max_files and len(out) >= max_files:
            break
    return out


@dataclass
class RuntimeConfig:
    max_turns: int
    max_tool_calls: int
    max_tokens: int
    obs_limit: int
    temperature: float
    thinking: bool
    save_transcripts: bool
    save_reasoning: bool
    ghidra_script_dir: Path


class StaticTools:
    def __init__(self, staged_path: Path, fmt: str, obs_limit: int, ghidra_script_dir: Path):
        self.staged_path = staged_path
        self.format = fmt
        self.obs_limit = obs_limit
        self.ghidra_script_dir = ghidra_script_dir
        self.used_tools: list[str] = []
        self.file_metadata: str | None = None
        self.mingw_objdump = shutil.which("x86_64-w64-mingw32-objdump") or "objdump"
        self.mingw_nm = shutil.which("x86_64-w64-mingw32-nm") or "nm"

    def _limit(self, text: str, limit: int) -> str:
        limit = min(limit, self.obs_limit)
        text = text.replace(str(self.staged_path), "<sample>")
        if len(text) > limit:
            return text[:limit] + f"\n[truncated to {limit} chars]"
        return text

    def _run(self, cmd: list[str], tool_name: str, limit: int, timeout: int = 45) -> str:
        self.used_tools.append(tool_name)
        try:
            proc = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                errors="replace",
                timeout=timeout,
                cwd=str(self.staged_path.parent),
                check=False,
            )
        except FileNotFoundError:
            return f"tool unavailable: {cmd[0]} not found"
        except subprocess.TimeoutExpired:
            return f"tool timeout after {timeout}s: {' '.join(cmd[:3])}"
        out = proc.stdout
        if proc.stderr:
            out += "\n[stderr]\n" + proc.stderr
        if not out.strip():
            out = f"{tool_name} completed with no output, exit={proc.returncode}"
        return self._limit(out, limit)

    def file(self, **_: Any) -> str:
        out = self._run(["file", "-b", str(self.staged_path)], "file", TOOL_LIMITS["file"], timeout=20)
        self.file_metadata = out
        return out

    def strings(self, min_length: Any = 5, **_: Any) -> str:
        try:
            n = int(min_length)
        except Exception:
            n = 5
        n = max(3, min(n, 32))
        return self._run(["strings", "-a", "-n", str(n), str(self.staged_path)], "strings", TOOL_LIMITS["strings"], 35)

    def readelf(self, **_: Any) -> str:
        if self.format != "ELF":
            return "readelf is only useful for ELF samples; use PE tools for PE files."
        return self._run(["readelf", "-aW", str(self.staged_path)], "readelf", TOOL_LIMITS["readelf"], 35)

    def objdump(self, **_: Any) -> str:
        if self.format == "PE":
            return self.pe_disasm()
        tool = shutil.which("x86_64-linux-gnu-objdump") or "objdump"
        cmd = [tool, "-d", "-M", "intel", str(self.staged_path)] if self._is_x86() else [tool, "-d", str(self.staged_path)]
        return self._run(cmd, "objdump", TOOL_LIMITS["objdump"], 60)

    def nm(self, **_: Any) -> str:
        nm_tool = self.mingw_nm if self.format == "PE" else "nm"
        return self._run([nm_tool, "-an", str(self.staged_path)], "nm", TOOL_LIMITS["nm"], 30)

    def disasm_func(self, function: Any = "", **_: Any) -> str:
        name = str(function or "").strip()
        if not name:
            self.used_tools.append("disasm_func")
            return "disasm_func requires a function symbol name."
        if self.format == "PE":
            cmd = [self.mingw_objdump, "-d", "-M", "intel", f"--disassemble={name}", str(self.staged_path)]
        else:
            tool = shutil.which("x86_64-linux-gnu-objdump") or "objdump"
            cmd = [tool, "-d", "-M", "intel", f"--disassemble={name}", str(self.staged_path)] if self._is_x86() else [tool, "-d", f"--disassemble={name}", str(self.staged_path)]
        return self._run(cmd, "disasm_func", TOOL_LIMITS["disasm_func"], 60)

    def pe_headers(self, **_: Any) -> str:
        if self.format != "PE":
            return "pe_headers is only useful for PE samples; use readelf for ELF files."
        a = self._run([self.mingw_objdump, "-f", str(self.staged_path)], "pe_headers", TOOL_LIMITS["pe_headers"], 25)
        b = self._run([self.mingw_objdump, "-p", str(self.staged_path)], "pe_headers", TOOL_LIMITS["pe_headers"], 45)
        return self._limit(a + "\n\n== private headers (-p) ==\n" + b, TOOL_LIMITS["pe_headers"])

    def pe_sections(self, **_: Any) -> str:
        if self.format != "PE":
            return "pe_sections is only useful for PE samples."
        hdr = self._run([self.mingw_objdump, "-h", str(self.staged_path)], "pe_sections", TOOL_LIMITS["pe_sections"], 25)
        return self._limit(hdr + "\n\n== per-section entropy ==\n" + self._section_entropy(hdr), TOOL_LIMITS["pe_sections"])

    def pe_imports(self, **_: Any) -> str:
        if self.format != "PE":
            return "pe_imports is only useful for PE samples."
        return self._run([self.mingw_objdump, "-x", str(self.staged_path)], "pe_imports", TOOL_LIMITS["pe_imports"], 45)

    def pe_exports(self, **_: Any) -> str:
        if self.format != "PE":
            return "pe_exports is only useful for PE samples."
        return self._run([self.mingw_objdump, "-p", str(self.staged_path)], "pe_exports", TOOL_LIMITS["pe_exports"], 45)

    def pe_disasm(self, **_: Any) -> str:
        if self.format != "PE":
            return "pe_disasm is only useful for PE samples; use objdump for ELF files."
        return self._run([self.mingw_objdump, "-d", "-M", "intel", str(self.staged_path)], "pe_disasm", TOOL_LIMITS["pe_disasm"], 70)

    def pe_symbols(self, **_: Any) -> str:
        if self.format != "PE":
            return "pe_symbols is only useful for PE samples; use nm for ELF files."
        return self._run([self.mingw_nm, str(self.staged_path)], "pe_symbols", TOOL_LIMITS["pe_symbols"], 30)

    def hexdump(self, length: Any = 1024, offset: Any = 0, **_: Any) -> str:
        length_i, offset_i = _dump_bounds(length, offset)
        if shutil.which("hexdump"):
            return self._run(["hexdump", "-C", "-n", str(length_i), "-s", str(offset_i), str(self.staged_path)], "hexdump", TOOL_LIMITS["hexdump"], 20)
        return self.xxd(length=length_i, offset=offset_i)

    def xxd(self, length: Any = 1024, offset: Any = 0, **_: Any) -> str:
        length_i, offset_i = _dump_bounds(length, offset)
        return self._run(["xxd", "-g", "1", "-l", str(length_i), "-s", str(offset_i), str(self.staged_path)], "xxd", TOOL_LIMITS["xxd"], 20)

    def entropy(self, **_: Any) -> str:
        self.used_tools.append("entropy")
        counts = [0] * 256
        total = 0
        with self.staged_path.open("rb") as fh:
            while True:
                chunk = fh.read(1 << 20)
                if not chunk:
                    break
                total += len(chunk)
                for b in chunk:
                    counts[b] += 1
        ent = 0.0 if total == 0 else -sum((c / total) * math.log2(c / total) for c in counts if c)
        return json.dumps({"bytes": total, "shannon_entropy": round(ent, 4)})

    def ghidra_summary(self, timeout: Any = 180, **_: Any) -> str:
        self.used_tools.append("ghidra_summary")
        try:
            timeout_i = max(30, min(int(timeout), 900))
        except Exception:
            timeout_i = 180
        headless = find_ghidra_headless()
        if not headless:
            return "ghidra unavailable: analyzeHeadless not found. Set GHIDRA_HEADLESS or add analyzeHeadless to PATH."
        script_dir = self.ghidra_script_dir
        script_file = script_dir / "AgentRESummary.java"
        if not script_file.exists():
            return f"ghidra script unavailable: {script_file}"
        with tempfile.TemporaryDirectory(prefix="agentre_ghidra_") as tmp:
            tmp_path = Path(tmp)
            out_path = tmp_path / "agentre_ghidra_summary.txt"
            cmd = [
                headless,
                str(tmp_path),
                "agentre_project",
                "-import",
                str(self.staged_path),
                "-overwrite",
                "-analysisTimeoutPerFile",
                str(timeout_i),
                "-scriptPath",
                str(script_dir),
                "-postScript",
                "AgentRESummary.java",
                str(out_path),
            ]
            try:
                proc = subprocess.run(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True,
                    errors="replace",
                    timeout=timeout_i + 90,
                    check=False,
                )
            except subprocess.TimeoutExpired:
                return f"ghidra timeout after {timeout_i + 90}s"
            out = out_path.read_text(encoding="utf-8", errors="replace") if out_path.exists() else ""
            log = proc.stdout + (("\n[stderr]\n" + proc.stderr) if proc.stderr else "")
            if out.strip():
                return self._limit(out + "\n\n== Ghidra log tail ==\n" + log[-4000:], TOOL_LIMITS["ghidra_summary"])
            return self._limit("Ghidra produced no summary file.\n\n" + log, TOOL_LIMITS["ghidra_summary"])

    def _is_x86(self) -> bool:
        meta = self.file_metadata or self.file()
        low = meta.lower()
        return "x86-64" in low or "80386" in low or "intel" in low or "amd64" in low

    def _section_entropy(self, section_table: str) -> str:
        try:
            data = self.staged_path.read_bytes()
        except Exception as exc:
            return f"<cannot read: {exc}>"
        rows: list[str] = []
        for line in section_table.splitlines():
            parts = line.split()
            if len(parts) >= 6 and parts[0].isdigit():
                try:
                    size, off = int(parts[2], 16), int(parts[5], 16)
                except ValueError:
                    continue
                blob = data[off : off + size]
                if not blob:
                    rows.append(f"{parts[1]:14s} size={size:<8} entropy=n/a")
                    continue
                counts = [0] * 256
                for b in blob:
                    counts[b] += 1
                n = len(blob)
                ent = -sum((v / n) * math.log2(v / n) for v in counts if v)
                flag = "  <HIGH: packed/encrypted?>" if ent > 7.2 else ""
                rows.append(f"{parts[1]:14s} size={size:<8} entropy={ent:4.2f}{flag}")
        return "\n".join(rows) or "<no section table>"


def find_ghidra_headless() -> str | None:
    env = os.environ.get("GHIDRA_HEADLESS")
    if env and Path(env).exists():
        return env
    found = shutil.which("analyzeHeadless")
    if found:
        return found

    # Keep this bounded. A recursive home-directory scan can hang on large
    # workstations, so only check common Ghidra install layouts.
    candidates = [
        Path("/opt/ghidra/support/analyzeHeadless"),
        Path("/usr/share/ghidra/support/analyzeHeadless"),
        Path.home() / "ghidra" / "support" / "analyzeHeadless",
    ]
    candidates.extend(Path("/opt").glob("ghidra*/support/analyzeHeadless") if Path("/opt").exists() else [])
    for path in candidates:
        if path.is_file() and os.access(path, os.X_OK):
            return str(path)
    return None


class LocalQwenBackend:
    def __init__(self, base_model: str, adapter: str | None, dtype: str, device_map: str):
        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer
        except Exception as exc:
            raise RuntimeError("Transformers local backend requires torch and transformers") from exc

        self.torch = torch
        tokenizer_path = adapter if adapter and Path(adapter).exists() else base_model
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
        torch_dtype = "auto" if dtype == "auto" else getattr(torch, dtype)
        self.model = AutoModelForCausalLM.from_pretrained(
            base_model,
            trust_remote_code=True,
            torch_dtype=torch_dtype,
            device_map=device_map,
        )
        if adapter:
            try:
                from peft import PeftModel
            except Exception as exc:
                raise RuntimeError("Loading a LoRA adapter requires peft") from exc
            self.model = PeftModel.from_pretrained(self.model, adapter)
        self.model.eval()

    def generate(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        max_new_tokens: int,
        temperature: float,
        thinking: bool,
    ) -> tuple[str, dict[str, int]]:
        prompt = self.tokenizer.apply_chat_template(
            messages,
            tools=tools,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=thinking,
        )
        inputs = self.tokenizer([prompt], return_tensors="pt")
        device = next(self.model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        do_sample = temperature > 0
        with self.torch.inference_mode():
            output = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=do_sample,
                temperature=temperature if do_sample else None,
                pad_token_id=self.tokenizer.eos_token_id,
            )
        input_len = int(inputs["input_ids"].shape[1])
        generated_ids = output[0][input_len:]
        text = self.tokenizer.decode(generated_ids, skip_special_tokens=False)
        if "<|im_end|>" in text:
            text = text.split("<|im_end|>", 1)[0]
        usage = {
            "prompt_tokens": input_len,
            "completion_tokens": int(generated_ids.shape[0]),
            "total_tokens": input_len + int(generated_ids.shape[0]),
        }
        return text, usage



def clean_llama_completion_output(text: str) -> str:
    lines: list[str] = []
    for line in text.splitlines():
        stripped = line.strip()
        if line.startswith("ggml_cuda_init:") or line.startswith("  Device "):
            continue
        if stripped in {"[end of text]", ""}:
            if lines:
                lines.append("")
            continue
        lines.append(line)
    cleaned = "\n".join(lines).replace(" [end of text]", "").replace("[end of text]", "").strip()
    return cleaned


class LlamaCppBackend:
    def __init__(
        self,
        model_path: str,
        llama_cli: str,
        ctx_size: int,
        threads: int,
        gpu_layers: int,
        llama_lora: str | None,
        extra_args: list[str],
    ):
        resolved_cli = shutil.which(llama_cli) or llama_cli
        cli_path = Path(resolved_cli)
        if cli_path.name == "llama-cli":
            sibling_completion = cli_path.with_name("llama-completion")
            if sibling_completion.exists():
                resolved_cli = str(sibling_completion)
        if not Path(resolved_cli).exists() and shutil.which(resolved_cli) is None:
            raise RuntimeError(f"llama.cpp completion binary not found: {llama_cli}. Set --llama-cli or AGENTRE_LLAMA_CLI.")
        if not Path(model_path).exists():
            raise RuntimeError(f"GGUF model not found: {model_path}")
        self.llama_cli = resolved_cli
        self.model_path = model_path
        self.ctx_size = ctx_size
        self.threads = threads
        self.gpu_layers = gpu_layers
        self.llama_lora = llama_lora
        self.extra_args = extra_args

    def generate(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        max_new_tokens: int,
        temperature: float,
        thinking: bool,
    ) -> tuple[str, dict[str, int]]:
        prompt = render_qwen_xml_chat(messages, tools, thinking=thinking, add_generation_prompt=True)

        def build_cmd(ctx_size: int) -> list[str]:
            cmd = [
                self.llama_cli,
                "-m",
                self.model_path,
                "-p",
                prompt,
                "-n",
                str(max_new_tokens),
                "-c",
                str(ctx_size),
                "--temp",
                str(max(0.0, temperature)),
                "--no-display-prompt",
                "-no-cnv",
                "--simple-io",
                "--no-warmup",
                "-lv",
                "1",
            ]
            if self.threads > 0:
                cmd.extend(["-t", str(self.threads)])
            if self.gpu_layers >= 0:
                cmd.extend(["-ngl", str(self.gpu_layers)])
            if self.llama_lora:
                cmd.extend(["--lora", self.llama_lora])
            cmd.extend(self.extra_args)
            return cmd

        ctx_size = self.ctx_size
        proc: subprocess.CompletedProcess[str] | None = None
        for attempt in range(2):
            try:
                proc = subprocess.run(
                    build_cmd(ctx_size),
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True,
                    errors="replace",
                    timeout=900,
                    check=False,
                )
            except FileNotFoundError as exc:
                raise RuntimeError(f"llama.cpp CLI not found: {self.llama_cli}") from exc

            if proc.returncode == 0:
                self.ctx_size = ctx_size
                break

            error_text = proc.stderr[-4000:] or proc.stdout[-4000:]
            match = re.search(r"prompt is too long \((\d+) tokens, max (\d+)\)", error_text)
            if match and attempt == 0:
                prompt_tokens = int(match.group(1))
                needed = prompt_tokens + max_new_tokens + 128
                next_ctx = max(ctx_size * 2, ((needed + 4095) // 4096) * 4096)
                print(
                    f"[xref 9b] llama.cpp prompt exceeded ctx={ctx_size}; retrying with ctx={next_ctx}",
                    file=sys.stderr,
                    flush=True,
                )
                ctx_size = next_ctx
                continue

            raise RuntimeError(
                "llama.cpp generation failed "
                f"exit={proc.returncode}: {error_text}"
            )

        if proc is None:
            raise RuntimeError("llama.cpp generation did not start")
        if proc.returncode != 0:
            error_text = proc.stderr[-4000:] or proc.stdout[-4000:]
            raise RuntimeError(
                "llama.cpp generation failed "
                f"exit={proc.returncode}: {error_text}"
            )

        text = clean_llama_completion_output(proc.stdout)
        if "<|im_end|>" in text:
            text = text.split("<|im_end|>", 1)[0]
        return text, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}


def build_messages(spec: str, sample_id: str, fmt: str, user_request: str) -> list[dict[str, Any]]:
    system = (
        spec.strip()
        + "\n\nRuntime rules:\n"
        + f"- You are analyzing staged sample `{sample_id}` only.\n"
        + f"- Detected format: {fmt}.\n"
        + "- The original path is hidden from the model to avoid path/name bias.\n"
        + "- Any tool path argument is ignored; tools are routed to the staged sample.\n"
        + "- Use static analysis only. Never ask to execute the sample.\n"
        + "- If evidence is insufficient, classify as unknown and state that closer disassembly with Ghidra or another disassembler is needed.\n"
        + "- For stripped, static, packed, encrypted, or sparse-string samples, prefer entropy, objdump/PE disassembly, and ghidra_summary before final_answer.\n"
        + "- Think internally, but do not include chain-of-thought in the final answer.\n"
    )
    user = (
        f"{user_request.strip()}\n\n"
        "Use the available reverse-engineering tools. Decide whether this sample is benign, malicious, hackware, or unknown. "
        "When finished, call final_answer with the structured JSON verdict."
    )
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


def make_history_tool_call(tc: dict[str, Any], fallback_id: str) -> dict[str, Any]:
    fn = tc.get("function", {})
    args = _parse_args(fn.get("arguments"))
    return {
        "id": tc.get("id") or fallback_id,
        "type": "function",
        "function": {
            "name": str(fn.get("name") or ""),
            "arguments": args,
        },
    }


def analyze_one(
    backend: Any,
    source_path: Path,
    sample_id: str,
    run_dir: Path,
    spec: str,
    cfg: RuntimeConfig,
    user_request: str,
) -> dict[str, Any]:
    fmt = detect_format(source_path)
    stage_dir = run_dir / "staged"
    stage_dir.mkdir(parents=True, exist_ok=True)
    suffix = ".elf" if fmt == "ELF" else ".exe" if fmt == "PE" else source_path.suffix
    staged_path = (stage_dir / f"{sample_id}{suffix}").resolve()
    shutil.copy2(source_path, staged_path)

    tools = StaticTools(staged_path, fmt, cfg.obs_limit, cfg.ghidra_script_dir)
    messages = build_messages(spec, sample_id, fmt, user_request)
    answer: dict[str, Any] | None = None
    error: str | None = None
    usage: dict[str, int] = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
    tool_calls_seen = 0
    transcript: list[dict[str, Any]] = list(messages)

    try:
        for turn in range(cfg.max_turns):
            if tool_calls_seen >= cfg.max_tool_calls:
                messages.append({"role": "user", "content": "Tool budget reached. Call final_answer now with your best JSON verdict."})
                transcript.append(messages[-1])

            generated, step_usage = backend.generate(
                messages,
                tool_definitions(),
                cfg.max_tokens,
                cfg.temperature,
                cfg.thinking,
            )
            for key, value in step_usage.items():
                usage[key] = usage.get(key, 0) + int(value)

            reasoning, visible = split_thinking(generated)
            tool_calls = parse_tool_calls(generated)
            assistant_content = strip_tool_xml(visible)
            assistant_msg: dict[str, Any] = {"role": "assistant", "content": assistant_content}
            if reasoning and cfg.save_reasoning:
                assistant_msg["reasoning_content"] = reasoning
            elif reasoning:
                assistant_msg["reasoning_content"] = ""
            if tool_calls:
                assistant_msg["tool_calls"] = [
                    make_history_tool_call(tc, f"call_{turn}_{idx}")
                    for idx, tc in enumerate(tool_calls)
                ]
            messages.append(assistant_msg)
            transcript.append(dict(assistant_msg))

            if not tool_calls:
                parsed = final_answer_from_args({"answer_json": assistant_content or generated})
                if parsed:
                    answer = parsed
                    break
                if turn < cfg.max_turns - 1:
                    prompt = "Call final_answer now with your compact JSON verdict."
                    messages.append({"role": "user", "content": prompt})
                    transcript.append(messages[-1])
                    continue
                break

            stop = False
            for idx, tc in enumerate(tool_calls):
                fn = tc.get("function", {})
                name = str(fn.get("name") or "")
                args = _parse_args(fn.get("arguments"))
                tcid = tc.get("id") or f"call_{turn}_{idx}"
                if name in FINAL_TOOLS:
                    answer = final_answer_from_args(args)
                    msg = {"role": "tool", "tool_call_id": tcid, "name": name, "content": "submitted"}
                    messages.append(msg)
                    transcript.append(msg)
                    stop = True
                    break
                canonical = TOOL_ALIASES.get(name)
                if canonical is None:
                    content = f"unknown tool {name}. Use the tools listed in the unified spec and final_answer."
                elif tool_calls_seen >= cfg.max_tool_calls:
                    content = "tool budget exhausted; submit final_answer."
                else:
                    tool_calls_seen += 1
                    try:
                        content = getattr(tools, canonical)(**args)
                    except Exception as exc:
                        content = f"tool error: {type(exc).__name__}: {exc}"
                msg = {
                    "role": "tool",
                    "tool_call_id": tcid,
                    "name": name,
                    "content": str(content)[: cfg.obs_limit],
                }
                messages.append(msg)
                transcript.append(msg)
            if stop:
                break
    except Exception as exc:
        error = f"{type(exc).__name__}: {exc}"

    prediction = normalize_prediction(answer)
    result = {
        "sample_id": sample_id,
        "source_path": str(source_path),
        "staged_path": str(staged_path),
        "format": fmt,
        "prediction": prediction,
        "answer": answer,
        "error": error,
        "tools_used": tools.used_tools,
        "tool_calls": tool_calls_seen,
        "usage": usage,
    }
    if cfg.save_transcripts:
        result["transcript"] = transcript
    return result


def write_outputs(results: list[dict[str, Any]], run_dir: Path) -> None:
    results_path = run_dir / "results.jsonl"
    with results_path.open("w", encoding="utf-8") as fh:
        for row in results:
            fh.write(json.dumps(row, ensure_ascii=True) + "\n")

    counts: dict[str, int] = {}
    by_format: dict[str, dict[str, int]] = {}
    for row in results:
        pred = str(row.get("prediction", "unknown"))
        fmt = str(row.get("format", "unknown"))
        counts[pred] = counts.get(pred, 0) + 1
        by_format.setdefault(fmt, {})
        by_format[fmt][pred] = by_format[fmt].get(pred, 0) + 1
    summary = {
        "total": len(results),
        "predictions": counts,
        "by_format": by_format,
        "errors": sum(1 for row in results if row.get("error")),
        "completed_at": datetime.now(timezone.utc).isoformat(),
    }
    (run_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")

    lines = [
        "# AgentRE Triage Summary",
        "",
        f"Total samples: {len(results)}",
        f"Errors: {summary['errors']}",
        "",
        "## Predictions",
        "",
    ]
    for label, count in sorted(counts.items()):
        lines.append(f"- {label}: {count}")
    lines += ["", "## Samples", ""]
    for row in results:
        lines.append(
            f"- `{row['sample_id']}` `{row['format']}` `{row['prediction']}` "
            f"tools={row['tool_calls']} path=`{row['source_path']}`"
        )
        answer = row.get("answer")
        if isinstance(answer, dict) and answer.get("summary"):
            lines.append(f"  - {answer['summary']}")
        if row.get("error"):
            lines.append(f"  - error: {row['error']}")
    (run_dir / "SUMMARY.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description="Analyze PE/ELF files with the AgentRE model and static tools.")
    parser.add_argument("target", help="File or directory to analyze.")
    parser.add_argument("--spec", default=str(DEFAULT_SPEC), help="Unified reverse-engineering spec.")
    parser.add_argument(
        "--backend",
        choices=["auto", "transformers", "llama-cpp"],
        default="auto",
        help="Inference backend. auto selects llama-cpp for .gguf models, otherwise transformers.",
    )
    parser.add_argument("--model", default=DEFAULT_BASE_MODEL, help="HF base model path/name or GGUF path for llama-cpp.")
    parser.add_argument("--adapter", default=DEFAULT_ADAPTER, help="Optional HF LoRA adapter path for transformers. Use '' to disable.")
    parser.add_argument("--output-dir", default="", help="Output directory. Defaults to runs/agentre_triage_<timestamp>.")
    parser.add_argument("--max-files", type=int, default=0, help="Directory mode limit; 0 means no limit.")
    parser.add_argument("--include-unknown", action="store_true", help="Include files that are not PE/ELF by magic bytes.")
    parser.add_argument("--max-turns", type=int, default=14)
    parser.add_argument("--max-tool-calls", type=int, default=12)
    parser.add_argument("--max-tokens", type=int, default=1200)
    parser.add_argument("--obs-limit", type=int, default=5000)
    parser.add_argument("--temperature", type=float, default=0.1)
    parser.add_argument("--dtype", default="auto", help="auto, bfloat16, float16, float32.")
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--llama-cli", default=DEFAULT_LLAMA_CLI, help="Path to llama.cpp llama-cli.")
    parser.add_argument("--ctx-size", type=int, default=32768, help="llama.cpp context size.")
    parser.add_argument("--threads", type=int, default=0, help="llama.cpp CPU threads; 0 lets llama.cpp choose.")
    parser.add_argument("--gpu-layers", type=int, default=-1, help="llama.cpp GPU layers; -1 leaves default.")
    parser.add_argument("--llama-lora", default="", help="Optional llama.cpp-compatible LoRA adapter. HF PEFT LoRA is not accepted here.")
    parser.add_argument("--llama-extra-arg", action="append", default=[], help="Extra raw argument passed to llama-cli. Repeatable.")
    parser.add_argument("--no-thinking", action="store_true", help="Disable Qwen thinking template flag.")
    parser.add_argument("--save-transcripts", action="store_true")
    parser.add_argument("--save-reasoning", action="store_true", help="Only meaningful with --save-transcripts.")
    parser.add_argument("--ghidra-script-dir", default=str(DEFAULT_GHIDRA_SCRIPT_DIR))
    parser.add_argument(
        "--request",
        default="Is this malicious or benign? Reverse engineer it and use the available tools.",
        help="User request inserted for each sample.",
    )
    parser.add_argument("--tool-smoke", action="store_true", help="Run static tools on the target and exit without loading the model.")
    args = parser.parse_args()

    target = Path(args.target).expanduser().resolve()
    targets = collect_targets(target, include_unknown=args.include_unknown, max_files=args.max_files)
    if not targets:
        print(f"No PE/ELF targets found under {target}", file=sys.stderr)
        return 2

    if args.output_dir:
        run_dir = Path(args.output_dir).expanduser().resolve()
    else:
        run_dir = (ROOT / "runs" / f"agentre_triage_{_now_stamp()}").resolve()
    run_dir.mkdir(parents=True, exist_ok=True)

    spec = Path(args.spec).read_text(encoding="utf-8")
    cfg = RuntimeConfig(
        max_turns=args.max_turns,
        max_tool_calls=args.max_tool_calls,
        max_tokens=args.max_tokens,
        obs_limit=args.obs_limit,
        temperature=args.temperature,
        thinking=not args.no_thinking,
        save_transcripts=args.save_transcripts,
        save_reasoning=args.save_reasoning,
        ghidra_script_dir=Path(args.ghidra_script_dir).expanduser().resolve(),
    )

    if args.tool_smoke:
        smoke_dir = run_dir / "tool_smoke"
        smoke_dir.mkdir(parents=True, exist_ok=True)
        for idx, path in enumerate(targets, 1):
            fmt = detect_format(path)
            staged = smoke_dir / f"sample_{idx:04d}{'.elf' if fmt == 'ELF' else '.exe' if fmt == 'PE' else path.suffix}"
            shutil.copy2(path, staged)
            tools = StaticTools(staged, fmt, cfg.obs_limit, cfg.ghidra_script_dir)
            print(f"== {path} ({fmt}) ==")
            print(tools.file())
            print(tools.strings()[:1200])
        return 0

    adapter = args.adapter.strip() or None
    backend_name = args.backend
    if backend_name == "auto":
        backend_name = "llama-cpp" if str(args.model).lower().endswith(".gguf") else "transformers"
    print(f"[xref 9b] backend={backend_name} model={args.model}", flush=True)
    if adapter and backend_name == "transformers":
        print(f"[xref 9b] adapter={adapter}", flush=True)
    elif adapter and backend_name == "llama-cpp":
        print("[xref 9b] note: --adapter is ignored by llama-cpp; use a merged GGUF or --llama-lora", flush=True)
    print(f"[xref 9b] thinking={'on' if cfg.thinking else 'off'} targets={len(targets)} output={run_dir}", flush=True)
    if backend_name == "transformers":
        backend = LocalQwenBackend(args.model, adapter, args.dtype, args.device_map)
    else:
        backend = LlamaCppBackend(
            args.model,
            args.llama_cli,
            args.ctx_size,
            args.threads,
            args.gpu_layers,
            args.llama_lora.strip() or None,
            args.llama_extra_arg,
        )

    results: list[dict[str, Any]] = []
    for idx, path in enumerate(targets, 1):
        sample_id = f"sample_{idx:04d}"
        print(f"[xref 9b] analyzing {sample_id} {path}", flush=True)
        result = analyze_one(backend, path, sample_id, run_dir, spec, cfg, args.request)
        results.append(result)
        print(
            f"[xref 9b] {sample_id} format={result['format']} prediction={result['prediction']} "
            f"tools={result['tool_calls']} error={result['error'] or ''}",
            flush=True,
        )
        write_outputs(results, run_dir)

    write_outputs(results, run_dir)
    print(f"[xref 9b] summary={run_dir / 'SUMMARY.md'}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())