mirror of
https://github.com/usestrix/strix.git
synced 2026-08-22 19:09:37 +02:00
473 lines
17 KiB
Python
473 lines
17 KiB
Python
import json
|
|
import logging
|
|
from collections.abc import Callable
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
from uuid import uuid4
|
|
|
|
from strix.telemetry import posthog
|
|
from strix.telemetry.scan_artifacts import ScanArtifactWriter
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_global_tracer: Optional["Tracer"] = None
|
|
|
|
|
|
def get_global_tracer() -> Optional["Tracer"]:
|
|
return _global_tracer
|
|
|
|
|
|
def set_global_tracer(tracer: "Tracer") -> None:
|
|
global _global_tracer # noqa: PLW0603
|
|
_global_tracer = tracer
|
|
|
|
|
|
class Tracer:
|
|
"""Per-scan in-memory state the TUI renders + per-scan artifact writer.
|
|
|
|
Holds live state the TUI reads (chat messages, agent tree, tool
|
|
executions, vulnerability reports, LLM usage). Writes vulnerability
|
|
markdown + CSV + final pentest report to ``strix_runs/<scan>/``.
|
|
|
|
Conversation history goes to the SDK's ``SQLiteSession`` instead;
|
|
SDK trace events are not persisted here.
|
|
"""
|
|
|
|
def __init__(self, run_name: str | None = None):
|
|
self.run_name = run_name
|
|
self.run_id = run_name or f"run-{uuid4().hex[:8]}"
|
|
self.start_time = datetime.now(UTC).isoformat()
|
|
self.end_time: str | None = None
|
|
|
|
self.agents: dict[str, dict[str, Any]] = {}
|
|
self.tool_executions: dict[int, dict[str, Any]] = {}
|
|
self.chat_messages: list[dict[str, Any]] = []
|
|
self._next_exec_id = 1
|
|
|
|
self.vulnerability_reports: list[dict[str, Any]] = []
|
|
self.final_scan_result: str | None = None
|
|
|
|
# LLM usage roll-up across all agents in this run.
|
|
self._llm_stats: dict[str, Any] = {
|
|
"input_tokens": 0,
|
|
"output_tokens": 0,
|
|
"cached_tokens": 0,
|
|
"cost": 0.0,
|
|
"requests": 0,
|
|
}
|
|
|
|
self.scan_results: dict[str, Any] | None = None
|
|
self.scan_config: dict[str, Any] | None = None
|
|
self.run_metadata: dict[str, Any] = {
|
|
"run_id": self.run_id,
|
|
"run_name": self.run_name,
|
|
"start_time": self.start_time,
|
|
"end_time": None,
|
|
"targets": [],
|
|
"status": "running",
|
|
}
|
|
self._run_dir: Path | None = None
|
|
self._writer: ScanArtifactWriter | None = None
|
|
self._next_message_id = 1
|
|
|
|
self.caido_url: str | None = None
|
|
self.vulnerability_found_callback: Callable[[dict[str, Any]], None] | None = None
|
|
|
|
def set_run_name(self, run_name: str) -> None:
|
|
self.run_name = run_name
|
|
self.run_id = run_name
|
|
self.run_metadata["run_name"] = run_name
|
|
self.run_metadata["run_id"] = run_name
|
|
self._run_dir = None
|
|
self._writer = None
|
|
|
|
def get_run_dir(self) -> Path:
|
|
if self._run_dir is None:
|
|
runs_dir = Path.cwd() / "strix_runs"
|
|
runs_dir.mkdir(exist_ok=True)
|
|
|
|
run_dir_name = self.run_name if self.run_name else self.run_id
|
|
self._run_dir = runs_dir / run_dir_name
|
|
self._run_dir.mkdir(exist_ok=True)
|
|
|
|
return self._run_dir
|
|
|
|
def hydrate_from_run_dir(self) -> None:
|
|
"""Reload prior-scan state from ``{run_dir}/`` for resume.
|
|
|
|
Called by :func:`run_strix_scan` before any new agent runs.
|
|
Restores:
|
|
|
|
- ``vulnerability_reports`` from ``vulnerabilities.json`` so
|
|
:meth:`add_vulnerability_report` doesn't allocate a colliding
|
|
``vuln-0001`` and overwrite the prior on-disk MD.
|
|
- ``run_metadata`` (start_time, run_id, targets, status) from
|
|
``run_metadata.json`` so audit-trail timestamps + the final
|
|
report's duration calc reflect the original scan, not just
|
|
this resume segment.
|
|
|
|
Idempotent on missing files (fresh runs land here too via the
|
|
same code path). **Raises on corruption** — silently swallowing
|
|
a corrupt ``vulnerabilities.json`` would let the next vuln
|
|
allocate ``vuln-0001`` and overwrite the prior MD on disk
|
|
(data loss). Caller is expected to fail the run loud and let
|
|
the user inspect ``{run_dir}`` or pick a fresh ``--run-name``.
|
|
"""
|
|
run_dir = self.get_run_dir()
|
|
|
|
meta_path = run_dir / "run_metadata.json"
|
|
if meta_path.exists():
|
|
try:
|
|
data = json.loads(meta_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise RuntimeError(
|
|
f"run_metadata.json at {meta_path} is unreadable: {exc}",
|
|
) from exc
|
|
if isinstance(data, dict):
|
|
if isinstance(data.get("start_time"), str):
|
|
self.start_time = data["start_time"]
|
|
self.run_metadata.update(
|
|
{
|
|
k: v
|
|
for k, v in data.items()
|
|
if k in {"run_id", "run_name", "start_time", "targets", "status"}
|
|
},
|
|
)
|
|
logger.info(
|
|
"tracer hydrated run_metadata from %s (start_time=%s)",
|
|
meta_path,
|
|
self.start_time,
|
|
)
|
|
|
|
json_path = run_dir / "vulnerabilities.json"
|
|
if json_path.exists():
|
|
try:
|
|
data = json.loads(json_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise RuntimeError(
|
|
f"vulnerabilities.json at {json_path} is corrupt ({exc}); "
|
|
f"refusing to start fresh — that would overwrite prior "
|
|
f"vulnerability MDs on disk. Inspect or delete the run dir.",
|
|
) from exc
|
|
if not isinstance(data, list):
|
|
raise RuntimeError(
|
|
f"vulnerabilities.json at {json_path} is not a list",
|
|
)
|
|
self.vulnerability_reports = [r for r in data if isinstance(r, dict)]
|
|
writer = self._get_writer()
|
|
for r in self.vulnerability_reports:
|
|
rid = r.get("id")
|
|
if isinstance(rid, str):
|
|
writer._saved_vuln_ids.add(rid)
|
|
logger.info(
|
|
"tracer hydrated %d vulnerability report(s) from %s",
|
|
len(self.vulnerability_reports),
|
|
json_path,
|
|
)
|
|
|
|
agents_path = run_dir / "agents.json"
|
|
if agents_path.exists():
|
|
try:
|
|
agents_data = json.loads(agents_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
# Caller will surface this same corruption via coordinator restore;
|
|
# no need to fail twice. Skip the agents/stats hydrate path.
|
|
agents_data = None
|
|
if isinstance(agents_data, dict):
|
|
self._hydrate_agents_tree(agents_data)
|
|
self._hydrate_llm_stats(agents_data)
|
|
|
|
def _hydrate_agents_tree(self, agents_data: dict[str, Any]) -> None:
|
|
"""Populate ``self.agents`` from the coordinator snapshot.
|
|
|
|
Without this, the TUI tree on resume would only show agents
|
|
currently running (mirrored by ``on_agent_start``); completed /
|
|
crashed / stopped children from the prior run would be invisible
|
|
even though the coordinator knows about them.
|
|
"""
|
|
statuses = agents_data.get("statuses") or {}
|
|
names = agents_data.get("names") or {}
|
|
parent_of = agents_data.get("parent_of") or {}
|
|
if not isinstance(statuses, dict):
|
|
return
|
|
timestamp = self.start_time
|
|
for agent_id, status in statuses.items():
|
|
if not isinstance(agent_id, str):
|
|
continue
|
|
self.agents[agent_id] = {
|
|
"id": agent_id,
|
|
"name": names.get(agent_id, agent_id) if isinstance(names, dict) else agent_id,
|
|
"parent_id": parent_of.get(agent_id) if isinstance(parent_of, dict) else None,
|
|
"status": status,
|
|
"created_at": timestamp,
|
|
"updated_at": timestamp,
|
|
}
|
|
logger.info("tracer hydrated %d agent(s) into tree", len(self.agents))
|
|
|
|
def _hydrate_llm_stats(self, agents_data: dict[str, Any]) -> None:
|
|
"""Seed ``self._llm_stats`` from the coordinator snapshot's counters.
|
|
|
|
This keeps the resumed scan's TUI footer cumulative instead of
|
|
resetting to zero.
|
|
"""
|
|
totals = {"input_tokens": 0, "output_tokens": 0, "cached_tokens": 0, "requests": 0}
|
|
bucket = agents_data.get("stats_live") or {}
|
|
if isinstance(bucket, dict):
|
|
for entry in bucket.values():
|
|
if isinstance(entry, dict):
|
|
totals["input_tokens"] += int(entry.get("in", 0) or 0)
|
|
totals["output_tokens"] += int(entry.get("out", 0) or 0)
|
|
totals["cached_tokens"] += int(entry.get("cached", 0) or 0)
|
|
totals["requests"] += int(entry.get("calls", 0) or 0)
|
|
for k, v in totals.items():
|
|
self._llm_stats[k] = v
|
|
logger.info(
|
|
"tracer hydrated llm stats from agents snapshot (in=%d out=%d cached=%d requests=%d)",
|
|
totals["input_tokens"],
|
|
totals["output_tokens"],
|
|
totals["cached_tokens"],
|
|
totals["requests"],
|
|
)
|
|
|
|
def _get_writer(self) -> ScanArtifactWriter:
|
|
if self._writer is None:
|
|
self._writer = ScanArtifactWriter(self.get_run_dir())
|
|
return self._writer
|
|
|
|
def add_vulnerability_report(
|
|
self,
|
|
title: str,
|
|
severity: str,
|
|
description: str | None = None,
|
|
impact: str | None = None,
|
|
target: str | None = None,
|
|
technical_analysis: str | None = None,
|
|
poc_description: str | None = None,
|
|
poc_script_code: str | None = None,
|
|
remediation_steps: str | None = None,
|
|
cvss: float | None = None,
|
|
cvss_breakdown: dict[str, str] | None = None,
|
|
endpoint: str | None = None,
|
|
method: str | None = None,
|
|
cve: str | None = None,
|
|
cwe: str | None = None,
|
|
code_locations: list[dict[str, Any]] | None = None,
|
|
) -> str:
|
|
report_id = f"vuln-{len(self.vulnerability_reports) + 1:04d}"
|
|
|
|
report: dict[str, Any] = {
|
|
"id": report_id,
|
|
"title": title.strip(),
|
|
"severity": severity.lower().strip(),
|
|
"timestamp": datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S UTC"),
|
|
}
|
|
|
|
if description:
|
|
report["description"] = description.strip()
|
|
if impact:
|
|
report["impact"] = impact.strip()
|
|
if target:
|
|
report["target"] = target.strip()
|
|
if technical_analysis:
|
|
report["technical_analysis"] = technical_analysis.strip()
|
|
if poc_description:
|
|
report["poc_description"] = poc_description.strip()
|
|
if poc_script_code:
|
|
report["poc_script_code"] = poc_script_code.strip()
|
|
if remediation_steps:
|
|
report["remediation_steps"] = remediation_steps.strip()
|
|
if cvss is not None:
|
|
report["cvss"] = cvss
|
|
if cvss_breakdown:
|
|
report["cvss_breakdown"] = cvss_breakdown
|
|
if endpoint:
|
|
report["endpoint"] = endpoint.strip()
|
|
if method:
|
|
report["method"] = method.strip()
|
|
if cve:
|
|
report["cve"] = cve.strip()
|
|
if cwe:
|
|
report["cwe"] = cwe.strip()
|
|
if code_locations:
|
|
report["code_locations"] = code_locations
|
|
|
|
self.vulnerability_reports.append(report)
|
|
logger.info(f"Added vulnerability report: {report_id} - {title}")
|
|
posthog.finding(severity)
|
|
|
|
if self.vulnerability_found_callback:
|
|
self.vulnerability_found_callback(report)
|
|
|
|
self.save_run_data()
|
|
return report_id
|
|
|
|
def get_existing_vulnerabilities(self) -> list[dict[str, Any]]:
|
|
return list(self.vulnerability_reports)
|
|
|
|
def update_scan_final_fields(
|
|
self,
|
|
executive_summary: str,
|
|
methodology: str,
|
|
technical_analysis: str,
|
|
recommendations: str,
|
|
) -> None:
|
|
self.scan_results = {
|
|
"scan_completed": True,
|
|
"executive_summary": executive_summary.strip(),
|
|
"methodology": methodology.strip(),
|
|
"technical_analysis": technical_analysis.strip(),
|
|
"recommendations": recommendations.strip(),
|
|
"success": True,
|
|
}
|
|
|
|
self.final_scan_result = f"""# Executive Summary
|
|
|
|
{executive_summary.strip()}
|
|
|
|
# Methodology
|
|
|
|
{methodology.strip()}
|
|
|
|
# Technical Analysis
|
|
|
|
{technical_analysis.strip()}
|
|
|
|
# Recommendations
|
|
|
|
{recommendations.strip()}
|
|
"""
|
|
|
|
logger.info("Updated scan final fields")
|
|
self.save_run_data(mark_complete=True)
|
|
posthog.end(self, exit_reason="finished_by_tool")
|
|
|
|
def log_chat_message(
|
|
self,
|
|
content: str,
|
|
role: str,
|
|
agent_id: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> int:
|
|
message_id = self._next_message_id
|
|
self._next_message_id += 1
|
|
|
|
self.chat_messages.append(
|
|
{
|
|
"message_id": message_id,
|
|
"content": content,
|
|
"role": role,
|
|
"agent_id": agent_id,
|
|
"timestamp": datetime.now(UTC).isoformat(),
|
|
"metadata": metadata or {},
|
|
}
|
|
)
|
|
return message_id
|
|
|
|
def set_scan_config(self, config: dict[str, Any]) -> None:
|
|
self.scan_config = config
|
|
self.run_metadata.update(
|
|
{
|
|
"targets": config.get("targets", []),
|
|
"user_instructions": config.get("user_instructions", ""),
|
|
"max_iterations": config.get("max_iterations", 200),
|
|
}
|
|
)
|
|
|
|
def save_run_data(self, mark_complete: bool = False) -> None:
|
|
if mark_complete:
|
|
if self.end_time is None:
|
|
self.end_time = datetime.now(UTC).isoformat()
|
|
self.run_metadata["end_time"] = self.end_time
|
|
self.run_metadata["status"] = "completed"
|
|
|
|
self._get_writer().save(
|
|
vulnerability_reports=self.vulnerability_reports,
|
|
final_scan_result=self.final_scan_result,
|
|
run_metadata=dict(self.run_metadata),
|
|
)
|
|
|
|
def log_tool_start(
|
|
self,
|
|
agent_id: str,
|
|
tool_name: str,
|
|
args: dict[str, Any] | None = None,
|
|
) -> int:
|
|
"""Record a tool invocation in flight. Returns an exec_id."""
|
|
exec_id = self._next_exec_id
|
|
self._next_exec_id += 1
|
|
self.tool_executions[exec_id] = {
|
|
"agent_id": agent_id,
|
|
"tool_name": tool_name,
|
|
"args": args or {},
|
|
"status": "running",
|
|
"result": None,
|
|
"timestamp": datetime.now(UTC).isoformat(),
|
|
}
|
|
return exec_id
|
|
|
|
def log_tool_end(self, agent_id: str, tool_name: str, result: Any) -> None:
|
|
"""Mark the most recent matching exec as completed."""
|
|
for exec_id in reversed(self.tool_executions):
|
|
entry = self.tool_executions[exec_id]
|
|
if (
|
|
entry.get("agent_id") == agent_id
|
|
and entry.get("tool_name") == tool_name
|
|
and entry.get("status") == "running"
|
|
):
|
|
entry["status"] = "completed"
|
|
entry["result"] = result
|
|
return
|
|
# No matching start (e.g. hooks added later in life) — record as completed.
|
|
exec_id = self._next_exec_id
|
|
self._next_exec_id += 1
|
|
self.tool_executions[exec_id] = {
|
|
"agent_id": agent_id,
|
|
"tool_name": tool_name,
|
|
"status": "completed",
|
|
"result": result,
|
|
"timestamp": datetime.now(UTC).isoformat(),
|
|
}
|
|
|
|
def get_real_tool_count(self) -> int:
|
|
return sum(
|
|
1
|
|
for exec_data in list(self.tool_executions.values())
|
|
if exec_data.get("tool_name") not in ["scan_start_info", "subagent_start_info"]
|
|
)
|
|
|
|
def get_total_llm_stats(self) -> dict[str, Any]:
|
|
"""Snapshot the run's aggregated LLM usage."""
|
|
stats = self._llm_stats
|
|
total = {
|
|
"input_tokens": int(stats["input_tokens"]),
|
|
"output_tokens": int(stats["output_tokens"]),
|
|
"cached_tokens": int(stats["cached_tokens"]),
|
|
"cost": round(float(stats["cost"]), 4),
|
|
"requests": int(stats["requests"]),
|
|
}
|
|
return {
|
|
"total": total,
|
|
"total_tokens": total["input_tokens"] + total["output_tokens"],
|
|
}
|
|
|
|
def record_llm_usage(
|
|
self,
|
|
*,
|
|
input_tokens: int = 0,
|
|
output_tokens: int = 0,
|
|
cached_tokens: int = 0,
|
|
cost: float = 0.0,
|
|
requests: int = 1,
|
|
) -> None:
|
|
"""Accumulate LLM usage from the orchestration hooks."""
|
|
self._llm_stats["input_tokens"] += input_tokens
|
|
self._llm_stats["output_tokens"] += output_tokens
|
|
self._llm_stats["cached_tokens"] += cached_tokens
|
|
self._llm_stats["cost"] += cost
|
|
self._llm_stats["requests"] += requests
|
|
|
|
def cleanup(self) -> None:
|
|
self.save_run_data(mark_complete=True)
|