#!/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"\s*\n]+)>\s*(.*?)\s*\s*",
re.DOTALL,
)
XML_PARAM_RE = re.compile(
r"\n]+)>\s*(.*?)\s*",
re.DOTALL,
)
JSON_TOOL_RE = re.compile(r"\s*(\{.*?\})\s*", re.DOTALL)
THINK_RE = re.compile(r"\s*(.*?)\s*\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 in the prompt, so the
# generated text may contain only the closing tag. Treat the prefix as
# hidden reasoning in that case.
end_tag = ""
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 = ["\n", f"\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"\n", value_text, "\n\n"])
parts.append("\n")
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")
for tool in tools:
chunks.append("\n")
chunks.append(json.dumps(tool, ensure_ascii=True))
chunks.append("\n")
chunks.append(
"\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n"
"\n\n"
"\nvalue_1\n\n"
"\nThis is the value for the second parameter\n"
"that can span\nmultiple lines\n\n"
"\n\n\n"
"\n"
"Reminder:\n"
"- Function calls MUST follow the specified format: an inner block must be nested within 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"
""
)
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("\n")
chunks.append(reasoning)
chunks.append("\n\n\n")
elif msg.get("reasoning_content") is not None:
chunks.append("\n\n\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\n")
chunks.append(str(tool_msg.get("content") or ""))
chunks.append("\n")
i += 1
chunks.append("<|im_end|>\n")
continue
i += 1
if add_generation_prompt:
chunks.append("<|im_start|>assistant\n")
if thinking:
chunks.append("\n")
else:
chunks.append("\n\n\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), "")
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""
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 = " " if ent > 7.2 else ""
rows.append(f"{parts[1]:14s} size={size:<8} entropy={ent:4.2f}{flag}")
return "\n".join(rows) or ""
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_.")
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())