first commit
This commit is contained in:
Zakaria
@@ -0,0 +1,12 @@
|
||||
"""Report/finding helpers."""
|
||||
|
||||
from strix.report.dedupe import check_duplicate
|
||||
from strix.report.state import ReportState, get_global_report_state, set_global_report_state
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ReportState",
|
||||
"check_duplicate",
|
||||
"get_global_report_state",
|
||||
"set_global_report_state",
|
||||
]
|
||||
@@ -0,0 +1,240 @@
|
||||
"""SDK-native vulnerability-report deduplication."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from agents.model_settings import ModelSettings
|
||||
from agents.models.interface import ModelTracing
|
||||
from openai.types.responses import ResponseOutputMessage
|
||||
|
||||
from strix.config import load_settings
|
||||
from strix.config.models import (
|
||||
DEFAULT_MODEL_RETRY,
|
||||
StrixProvider,
|
||||
configure_sdk_model_defaults,
|
||||
)
|
||||
from strix.report.state import get_global_report_state
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agents.items import ModelResponse
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEDUPE_SYSTEM_PROMPT = """You are an expert vulnerability report deduplication judge.
|
||||
Your task is to determine if a candidate vulnerability report describes the SAME vulnerability
|
||||
as any existing report.
|
||||
|
||||
CRITICAL DEDUPLICATION RULES:
|
||||
|
||||
1. SAME VULNERABILITY means:
|
||||
- Same root cause (e.g., "missing input validation" not just "SQL injection")
|
||||
- Same affected component/endpoint/file (exact match or clear overlap)
|
||||
- Same exploitation method or attack vector
|
||||
- Would be fixed by the same code change/patch
|
||||
|
||||
2. NOT DUPLICATES if:
|
||||
- Different endpoints even with same vulnerability type (e.g., SQLi in /login vs /search)
|
||||
- Different parameters in same endpoint (e.g., XSS in 'name' vs 'comment' field)
|
||||
- Different root causes (e.g., stored XSS vs reflected XSS in same field)
|
||||
- Different severity levels due to different impact
|
||||
- One is authenticated, other is unauthenticated
|
||||
|
||||
3. ARE DUPLICATES even if:
|
||||
- Titles are worded differently
|
||||
- Descriptions have different level of detail
|
||||
- PoC uses different payloads but exploits same issue
|
||||
- One report is more thorough than another
|
||||
- Minor variations in technical analysis
|
||||
|
||||
COMPARISON GUIDELINES:
|
||||
- Focus on the technical root cause, not surface-level similarities
|
||||
- Same vulnerability type (SQLi, XSS) doesn't mean duplicate - location matters
|
||||
- Consider the fix: would fixing one also fix the other?
|
||||
- When uncertain, lean towards NOT duplicate
|
||||
|
||||
FIELDS TO ANALYZE:
|
||||
- title, description: General vulnerability info
|
||||
- target, endpoint, method: Exact location of vulnerability
|
||||
- technical_analysis: Root cause details
|
||||
- poc_description: How it's exploited
|
||||
- impact: What damage it can cause
|
||||
|
||||
Respond with a single JSON object and nothing else:
|
||||
|
||||
{
|
||||
"is_duplicate": true,
|
||||
"duplicate_id": "vuln-0001",
|
||||
"confidence": 0.95,
|
||||
"reason": "Both reports describe SQL injection in /api/login via the username parameter"
|
||||
}
|
||||
|
||||
Or, if not a duplicate:
|
||||
|
||||
{
|
||||
"is_duplicate": false,
|
||||
"duplicate_id": "",
|
||||
"confidence": 0.90,
|
||||
"reason": "Different endpoints: candidate is /api/search, existing is /api/login"
|
||||
}
|
||||
|
||||
Rules:
|
||||
- ``is_duplicate`` is a boolean.
|
||||
- ``duplicate_id`` is the exact id from existing reports, or "" if not a duplicate.
|
||||
- ``confidence`` is a number between 0 and 1.
|
||||
- ``reason`` is a specific explanation mentioning endpoint/parameter/root cause.
|
||||
- Output ONLY the JSON object — no surrounding prose, no code fences."""
|
||||
|
||||
|
||||
def _prepare_report_for_comparison(report: dict[str, Any]) -> dict[str, Any]:
|
||||
relevant_fields = [
|
||||
"id",
|
||||
"title",
|
||||
"description",
|
||||
"impact",
|
||||
"target",
|
||||
"technical_analysis",
|
||||
"poc_description",
|
||||
"endpoint",
|
||||
"method",
|
||||
]
|
||||
|
||||
cleaned = {}
|
||||
for field in relevant_fields:
|
||||
if report.get(field):
|
||||
value = report[field]
|
||||
if isinstance(value, str) and len(value) > 8000:
|
||||
value = value[:8000] + "...[truncated]"
|
||||
cleaned[field] = value
|
||||
|
||||
return cleaned
|
||||
|
||||
|
||||
def _parse_dedupe_response(content: str) -> dict[str, Any]:
|
||||
text = content.strip()
|
||||
if text.startswith("```"):
|
||||
text = text.strip("`")
|
||||
if text.lower().startswith("json"):
|
||||
text = text[4:]
|
||||
text = text.strip()
|
||||
start = text.find("{")
|
||||
end = text.rfind("}")
|
||||
if start == -1 or end == -1 or end <= start:
|
||||
raise ValueError(f"No JSON object found in dedupe response: {content[:500]}")
|
||||
parsed = json.loads(text[start : end + 1])
|
||||
|
||||
duplicate_id = str(parsed.get("duplicate_id") or "")[:64]
|
||||
reason = str(parsed.get("reason") or "")[:500]
|
||||
try:
|
||||
confidence = float(parsed.get("confidence", 0.0))
|
||||
except (TypeError, ValueError):
|
||||
confidence = 0.0
|
||||
|
||||
return {
|
||||
"is_duplicate": bool(parsed.get("is_duplicate", False)),
|
||||
"duplicate_id": duplicate_id,
|
||||
"confidence": confidence,
|
||||
"reason": reason,
|
||||
}
|
||||
|
||||
|
||||
def _extract_text(response: ModelResponse) -> str:
|
||||
parts: list[str] = []
|
||||
for item in response.output:
|
||||
if not isinstance(item, ResponseOutputMessage):
|
||||
continue
|
||||
for chunk in item.content:
|
||||
text = getattr(chunk, "text", None)
|
||||
if text:
|
||||
parts.append(text)
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
async def check_duplicate(
|
||||
candidate: dict[str, Any], existing_reports: list[dict[str, Any]]
|
||||
) -> dict[str, Any]:
|
||||
if not existing_reports:
|
||||
return {
|
||||
"is_duplicate": False,
|
||||
"duplicate_id": "",
|
||||
"confidence": 1.0,
|
||||
"reason": "No existing reports to compare against",
|
||||
}
|
||||
|
||||
try:
|
||||
settings = load_settings()
|
||||
model_name = settings.llm.model
|
||||
if not model_name:
|
||||
return {
|
||||
"is_duplicate": False,
|
||||
"duplicate_id": "",
|
||||
"confidence": 0.0,
|
||||
"reason": "STRIX_LLM not configured; skipping dedupe check",
|
||||
}
|
||||
|
||||
candidate_cleaned = _prepare_report_for_comparison(candidate)
|
||||
existing_cleaned = [_prepare_report_for_comparison(r) for r in existing_reports]
|
||||
comparison_data = {"candidate": candidate_cleaned, "existing_reports": existing_cleaned}
|
||||
|
||||
user_msg = (
|
||||
f"Compare this candidate vulnerability against existing reports:\n\n"
|
||||
f"{json.dumps(comparison_data, indent=2)}\n\n"
|
||||
f"Respond with ONLY the JSON object described in the system prompt."
|
||||
)
|
||||
|
||||
configure_sdk_model_defaults(settings)
|
||||
resolved_model = model_name.strip()
|
||||
model = StrixProvider().get_model(resolved_model)
|
||||
response = await model.get_response(
|
||||
system_instructions=DEDUPE_SYSTEM_PROMPT,
|
||||
input=user_msg,
|
||||
model_settings=ModelSettings(retry=DEFAULT_MODEL_RETRY, include_usage=True),
|
||||
tools=[],
|
||||
output_schema=None,
|
||||
handoffs=[],
|
||||
tracing=ModelTracing.DISABLED,
|
||||
previous_response_id=None,
|
||||
conversation_id=None,
|
||||
prompt=None,
|
||||
)
|
||||
report_state = get_global_report_state()
|
||||
if report_state is not None:
|
||||
report_state.record_sdk_usage(
|
||||
agent_id="dedupe",
|
||||
agent_name="dedupe",
|
||||
model=resolved_model,
|
||||
usage=response.usage,
|
||||
)
|
||||
content = _extract_text(response)
|
||||
if not content:
|
||||
return {
|
||||
"is_duplicate": False,
|
||||
"duplicate_id": "",
|
||||
"confidence": 0.0,
|
||||
"reason": "Empty response from LLM",
|
||||
}
|
||||
|
||||
result = _parse_dedupe_response(content)
|
||||
|
||||
logger.info(
|
||||
"Deduplication check: is_duplicate=%s, confidence=%.2f, reason=%s",
|
||||
result["is_duplicate"],
|
||||
result["confidence"],
|
||||
result["reason"][:100],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Error during vulnerability deduplication check")
|
||||
return {
|
||||
"is_duplicate": False,
|
||||
"duplicate_id": "",
|
||||
"confidence": 0.0,
|
||||
"reason": f"Deduplication check failed: {e}",
|
||||
"error": str(e),
|
||||
}
|
||||
else:
|
||||
return result
|
||||
@@ -0,0 +1,394 @@
|
||||
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 agents.usage import Usage
|
||||
|
||||
from strix.core.paths import run_dir_for
|
||||
from strix.report.usage import LLMUsageLedger
|
||||
from strix.report.writer import (
|
||||
read_run_record,
|
||||
write_executive_report,
|
||||
write_run_record,
|
||||
write_vulnerabilities,
|
||||
)
|
||||
from strix.telemetry import posthog, scarf
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_global_report_state: Optional["ReportState"] = None
|
||||
|
||||
|
||||
def get_global_report_state() -> Optional["ReportState"]:
|
||||
return _global_report_state
|
||||
|
||||
|
||||
def set_global_report_state(report_state: "ReportState") -> None:
|
||||
global _global_report_state # noqa: PLW0603
|
||||
_global_report_state = report_state
|
||||
|
||||
|
||||
class ReportState:
|
||||
"""Per-scan product artifact state plus artifact writer.
|
||||
|
||||
The Agents SDK owns model/tool execution, tracing, and conversation
|
||||
persistence. This store keeps only Strix-owned scan artifacts and
|
||||
report metadata. Live UI projections belong to the interface layer.
|
||||
|
||||
It does not consume SDK tracing processors.
|
||||
"""
|
||||
|
||||
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.vulnerability_reports: list[dict[str, Any]] = []
|
||||
self.final_scan_result: str | None = None
|
||||
|
||||
self.scan_results: dict[str, Any] | None = None
|
||||
self.scan_config: dict[str, Any] | None = None
|
||||
self._llm_usage = LLMUsageLedger()
|
||||
self.run_record: dict[str, Any] = {
|
||||
"run_id": self.run_id,
|
||||
"run_name": self.run_name,
|
||||
"start_time": self.start_time,
|
||||
"end_time": None,
|
||||
"status": "running",
|
||||
"targets_info": [],
|
||||
"llm_usage": self._build_llm_usage_record(),
|
||||
}
|
||||
self._run_dir: Path | None = None
|
||||
self._saved_vuln_ids: set[str] = set()
|
||||
|
||||
self.caido_url: str | None = None
|
||||
self.vulnerability_found_callback: Callable[[dict[str, Any]], None] | None = None
|
||||
|
||||
def get_run_dir(self) -> Path:
|
||||
if self._run_dir is None:
|
||||
run_dir_name = self.run_name if self.run_name else self.run_id
|
||||
self._run_dir = run_dir_for(run_dir_name)
|
||||
self._run_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return self._run_dir
|
||||
|
||||
def hydrate_from_run_dir(self) -> None:
|
||||
"""Reload prior-scan state from ``{run_dir}/`` for resume.
|
||||
|
||||
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_record`` from ``run.json`` so timestamps, run inputs,
|
||||
status, and final report state have one public source of truth.
|
||||
|
||||
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()
|
||||
|
||||
data = read_run_record(run_dir)
|
||||
if data:
|
||||
self.run_record.update(data)
|
||||
if isinstance(data.get("start_time"), str):
|
||||
self.start_time = data["start_time"]
|
||||
if isinstance(data.get("end_time"), str):
|
||||
self.end_time = data["end_time"]
|
||||
scan_results = data.get("scan_results")
|
||||
if isinstance(scan_results, dict):
|
||||
self.scan_results = scan_results
|
||||
self.final_scan_result = self._format_final_scan_result(scan_results)
|
||||
self._hydrate_llm_usage(data.get("llm_usage"))
|
||||
logger.info("report state hydrated run.json from %s", run_dir)
|
||||
|
||||
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)]
|
||||
for r in self.vulnerability_reports:
|
||||
rid = r.get("id")
|
||||
if isinstance(rid, str):
|
||||
self._saved_vuln_ids.add(rid)
|
||||
logger.info(
|
||||
"report state hydrated %d vulnerability report(s)",
|
||||
len(self.vulnerability_reports),
|
||||
)
|
||||
|
||||
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,
|
||||
agent_id: str | None = None,
|
||||
agent_name: str | 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
|
||||
if agent_id:
|
||||
report["agent_id"] = agent_id
|
||||
if agent_name:
|
||||
report["agent_name"] = agent_name
|
||||
|
||||
self.vulnerability_reports.append(report)
|
||||
logger.info(f"Added vulnerability report: {report_id} - {title}")
|
||||
posthog.finding(severity)
|
||||
scarf.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 record_sdk_usage(
|
||||
self,
|
||||
*,
|
||||
agent_id: str,
|
||||
usage: Usage | None,
|
||||
agent_name: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> None:
|
||||
"""Record SDK-native token usage for one completed model run/cycle."""
|
||||
if self._llm_usage.record(
|
||||
agent_id=agent_id,
|
||||
agent_name=agent_name,
|
||||
model=model,
|
||||
usage=usage,
|
||||
):
|
||||
self.save_run_data()
|
||||
|
||||
def record_observed_llm_cost(self, cost: float) -> None:
|
||||
self._llm_usage.record_observed_cost(cost)
|
||||
|
||||
def get_total_llm_usage(self) -> dict[str, Any]:
|
||||
return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record())
|
||||
|
||||
def get_total_llm_cost(self) -> float:
|
||||
"""Live accumulated LLM cost, independent of the persisted run-record snapshot."""
|
||||
return self._llm_usage.total_cost
|
||||
|
||||
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 = self._format_final_scan_result(self.scan_results)
|
||||
self.run_record["scan_results"] = self.scan_results
|
||||
|
||||
logger.info("Updated scan final fields")
|
||||
self.save_run_data(mark_complete=True)
|
||||
posthog.end(self, exit_reason="finished_by_tool")
|
||||
scarf.end(self, exit_reason="finished_by_tool")
|
||||
|
||||
def set_scan_config(self, config: dict[str, Any]) -> None:
|
||||
self.scan_config = config
|
||||
self.run_record["status"] = "running"
|
||||
self.run_record["end_time"] = None
|
||||
self.run_record.pop("scan_results", None)
|
||||
self.end_time = None
|
||||
self.scan_results = None
|
||||
self.final_scan_result = None
|
||||
self.run_record.update(
|
||||
{
|
||||
"targets_info": config.get("targets", []),
|
||||
"instruction": config.get("user_instructions", ""),
|
||||
"scan_mode": config.get("scan_mode", "deep"),
|
||||
"diff_scope": config.get("diff_scope", {"active": False}),
|
||||
"non_interactive": bool(config.get("non_interactive", False)),
|
||||
"local_sources": config.get("local_sources", []),
|
||||
"scope_mode": config.get("scope_mode", "auto"),
|
||||
"diff_base": config.get("diff_base"),
|
||||
}
|
||||
)
|
||||
|
||||
def save_run_data(self, mark_complete: bool = False, status: str | None = None) -> None:
|
||||
if mark_complete:
|
||||
self.end_time = datetime.now(UTC).isoformat()
|
||||
self.run_record["end_time"] = self.end_time
|
||||
self.run_record["status"] = "completed"
|
||||
elif status and self.run_record.get("status") != "completed":
|
||||
current_status = self.run_record.get("status")
|
||||
if status == "stopped" and current_status in {"failed", "interrupted"}:
|
||||
status = str(current_status)
|
||||
if self.end_time is None:
|
||||
self.end_time = datetime.now(UTC).isoformat()
|
||||
self.run_record["end_time"] = self.end_time
|
||||
self.run_record["status"] = status
|
||||
|
||||
self._sync_llm_usage_record()
|
||||
self._save_artifacts()
|
||||
|
||||
def cleanup(self, status: str = "stopped") -> None:
|
||||
self.save_run_data(status=status)
|
||||
|
||||
def _format_final_scan_result(self, scan_results: dict[str, Any]) -> str:
|
||||
return f"""# Executive Summary
|
||||
|
||||
{str(scan_results.get("executive_summary", "")).strip()}
|
||||
|
||||
# Methodology
|
||||
|
||||
{str(scan_results.get("methodology", "")).strip()}
|
||||
|
||||
# Technical Analysis
|
||||
|
||||
{str(scan_results.get("technical_analysis", "")).strip()}
|
||||
|
||||
# Recommendations
|
||||
|
||||
{str(scan_results.get("recommendations", "")).strip()}
|
||||
"""
|
||||
|
||||
def _save_artifacts(self) -> None:
|
||||
"""Write scan artifacts under ``run_dir``."""
|
||||
run_dir = self.get_run_dir()
|
||||
try:
|
||||
run_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if self.final_scan_result:
|
||||
write_executive_report(run_dir, self.final_scan_result)
|
||||
|
||||
if self.vulnerability_reports:
|
||||
write_vulnerabilities(run_dir, self.vulnerability_reports, self._saved_vuln_ids)
|
||||
|
||||
write_run_record(run_dir, self.run_record)
|
||||
|
||||
logger.info("Essential scan data saved to: %s", run_dir)
|
||||
except (OSError, RuntimeError):
|
||||
logger.exception("Failed to save scan data")
|
||||
|
||||
def _sync_llm_usage_record(self) -> None:
|
||||
self.run_record["llm_usage"] = self._build_llm_usage_record()
|
||||
|
||||
def _build_llm_usage_record(self) -> dict[str, Any]:
|
||||
return self._llm_usage.to_record()
|
||||
|
||||
def _hydrate_llm_usage(self, raw_usage: Any) -> None:
|
||||
self._llm_usage.hydrate(raw_usage)
|
||||
self._sync_llm_usage_record()
|
||||
|
||||
|
||||
def litellm_cost_callback(
|
||||
kwargs: Any,
|
||||
completion_response: Any,
|
||||
_start_time: Any = None,
|
||||
_end_time: Any = None,
|
||||
) -> None:
|
||||
"""LiteLLM ``success_callback`` adapter; forwards observed cost to the active scan."""
|
||||
cost: float | None = None
|
||||
raw = kwargs.get("response_cost") if isinstance(kwargs, dict) else None
|
||||
if isinstance(raw, int | float) and raw > 0:
|
||||
cost = float(raw)
|
||||
|
||||
if cost is None:
|
||||
hidden = getattr(completion_response, "_hidden_params", None) or {}
|
||||
candidate = hidden.get("response_cost") if isinstance(hidden, dict) else None
|
||||
if isinstance(candidate, int | float) and candidate > 0:
|
||||
cost = float(candidate)
|
||||
else:
|
||||
headers = hidden.get("additional_headers") or {} if isinstance(hidden, dict) else {}
|
||||
raw = (
|
||||
headers.get("llm_provider-x-litellm-response-cost")
|
||||
if isinstance(headers, dict)
|
||||
else None
|
||||
)
|
||||
try:
|
||||
value = float(raw) if raw is not None else None
|
||||
except (TypeError, ValueError):
|
||||
value = None
|
||||
if value is not None and value > 0:
|
||||
cost = value
|
||||
|
||||
if cost is None or cost <= 0:
|
||||
return
|
||||
report_state = get_global_report_state()
|
||||
if report_state is None:
|
||||
return
|
||||
try:
|
||||
report_state.record_observed_llm_cost(cost)
|
||||
except Exception:
|
||||
logger.exception("Failed to record observed LiteLLM cost")
|
||||
@@ -0,0 +1,262 @@
|
||||
"""SDK-native LLM usage aggregation for scan reports."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from agents.usage import Usage, deserialize_usage, serialize_usage
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMUsageLedger:
|
||||
"""Aggregate SDK ``Usage`` objects and attach best-effort cost estimates."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._total_usage = Usage()
|
||||
self._agent_usage: dict[str, Usage] = {}
|
||||
self._agent_metadata: dict[str, dict[str, str]] = {}
|
||||
self._total_cost = 0.0
|
||||
|
||||
def record(
|
||||
self,
|
||||
*,
|
||||
agent_id: str,
|
||||
usage: Usage | None,
|
||||
agent_name: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> bool:
|
||||
if usage is None or not _usage_has_activity(usage):
|
||||
return False
|
||||
|
||||
normalized_agent_id = str(agent_id or "unknown")
|
||||
self._total_usage.add(usage)
|
||||
self._agent_usage.setdefault(normalized_agent_id, Usage()).add(usage)
|
||||
|
||||
metadata = self._agent_metadata.setdefault(normalized_agent_id, {})
|
||||
if agent_name:
|
||||
metadata["agent_name"] = agent_name
|
||||
if model:
|
||||
metadata["model"] = model
|
||||
|
||||
if not _is_litellm_routed(model):
|
||||
estimated = _estimate_litellm_cost(usage, model)
|
||||
if estimated:
|
||||
self._total_cost += estimated
|
||||
|
||||
return True
|
||||
|
||||
def record_observed_cost(self, cost: float) -> None:
|
||||
if isinstance(cost, int | float) and cost > 0:
|
||||
self._total_cost += float(cost)
|
||||
|
||||
@property
|
||||
def total_cost(self) -> float:
|
||||
return _round_cost(self._total_cost)
|
||||
|
||||
def to_record(self) -> dict[str, Any]:
|
||||
record = serialize_usage(self._total_usage)
|
||||
record["cost"] = _round_cost(self._total_cost)
|
||||
record["agents"] = []
|
||||
|
||||
agent_tokens = {aid: _resolve_total_tokens(u) for aid, u in self._agent_usage.items()}
|
||||
total_tokens = sum(agent_tokens.values())
|
||||
for agent_id in sorted(self._agent_usage):
|
||||
usage = self._agent_usage[agent_id]
|
||||
metadata = self._agent_metadata.get(agent_id, {})
|
||||
agent_cost = (
|
||||
self._total_cost * (agent_tokens[agent_id] / total_tokens) if total_tokens else 0.0
|
||||
)
|
||||
|
||||
agent_record = serialize_usage(usage)
|
||||
agent_record.update(
|
||||
{
|
||||
"agent_id": agent_id,
|
||||
"agent_name": metadata.get("agent_name") or agent_id,
|
||||
"model": metadata.get("model"),
|
||||
"cost": _round_cost(agent_cost),
|
||||
}
|
||||
)
|
||||
record["agents"].append(agent_record)
|
||||
|
||||
return record
|
||||
|
||||
def hydrate(self, raw_usage: Any) -> None:
|
||||
self._total_usage = Usage()
|
||||
self._agent_usage.clear()
|
||||
self._agent_metadata.clear()
|
||||
self._total_cost = 0.0
|
||||
|
||||
if not isinstance(raw_usage, dict):
|
||||
return
|
||||
|
||||
try:
|
||||
self._total_usage = deserialize_usage(raw_usage)
|
||||
except Exception:
|
||||
logger.exception("Failed to hydrate aggregate llm_usage from run.json")
|
||||
self._total_usage = Usage()
|
||||
|
||||
self._total_cost = _float_or_zero(raw_usage.get("cost"))
|
||||
|
||||
for raw_agent in raw_usage.get("agents") or []:
|
||||
if not isinstance(raw_agent, dict):
|
||||
continue
|
||||
agent_id = str(raw_agent.get("agent_id") or "").strip()
|
||||
if not agent_id:
|
||||
continue
|
||||
try:
|
||||
self._agent_usage[agent_id] = deserialize_usage(raw_agent)
|
||||
except Exception:
|
||||
logger.exception("Failed to hydrate llm_usage for agent %s", agent_id)
|
||||
self._agent_usage[agent_id] = Usage()
|
||||
|
||||
metadata: dict[str, str] = {}
|
||||
agent_name = raw_agent.get("agent_name")
|
||||
model = raw_agent.get("model")
|
||||
if isinstance(agent_name, str) and agent_name:
|
||||
metadata["agent_name"] = agent_name
|
||||
if isinstance(model, str) and model:
|
||||
metadata["model"] = model
|
||||
self._agent_metadata[agent_id] = metadata
|
||||
|
||||
|
||||
def _resolve_total_tokens(usage: Usage) -> int:
|
||||
total = max(0, int(usage.total_tokens or 0))
|
||||
if total > 0:
|
||||
return total
|
||||
prompt = _int_or_zero(getattr(usage, "input_tokens", 0))
|
||||
completion = _int_or_zero(getattr(usage, "output_tokens", 0))
|
||||
return prompt + completion
|
||||
|
||||
|
||||
def _is_litellm_routed(model: str | None) -> bool:
|
||||
if not model:
|
||||
return False
|
||||
name = model.strip().lower()
|
||||
if "/" not in name:
|
||||
return False
|
||||
return not name.startswith("openai/")
|
||||
|
||||
|
||||
def _usage_has_activity(usage: Usage) -> bool:
|
||||
return bool(
|
||||
usage.requests
|
||||
or usage.input_tokens
|
||||
or usage.output_tokens
|
||||
or usage.total_tokens
|
||||
or usage.request_usage_entries
|
||||
)
|
||||
|
||||
|
||||
def _estimate_litellm_cost(usage: Usage, model: str | None) -> float | None:
|
||||
litellm_model = _litellm_model_name(model)
|
||||
if not litellm_model:
|
||||
return None
|
||||
|
||||
entries = list(usage.request_usage_entries)
|
||||
if not entries:
|
||||
return _estimate_litellm_entry_cost(usage, litellm_model)
|
||||
|
||||
total = 0.0
|
||||
estimated_any = False
|
||||
for entry in entries:
|
||||
cost = _estimate_litellm_entry_cost(entry, litellm_model)
|
||||
if cost is None:
|
||||
continue
|
||||
total += cost
|
||||
estimated_any = True
|
||||
|
||||
return total if estimated_any else None
|
||||
|
||||
|
||||
def _estimate_litellm_entry_cost(entry: Any, model: str) -> float | None:
|
||||
prompt_tokens = _int_or_zero(getattr(entry, "input_tokens", 0))
|
||||
completion_tokens = _int_or_zero(getattr(entry, "output_tokens", 0))
|
||||
total_tokens = _int_or_zero(getattr(entry, "total_tokens", 0))
|
||||
if total_tokens <= 0:
|
||||
total_tokens = prompt_tokens + completion_tokens
|
||||
if total_tokens <= 0:
|
||||
return None
|
||||
|
||||
usage_payload: dict[str, Any] = {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
prompt_details = _details_to_dict(getattr(entry, "input_tokens_details", None))
|
||||
completion_details = _details_to_dict(getattr(entry, "output_tokens_details", None))
|
||||
if prompt_details:
|
||||
usage_payload["prompt_tokens_details"] = prompt_details
|
||||
if completion_details:
|
||||
usage_payload["completion_tokens_details"] = completion_details
|
||||
|
||||
from litellm import completion_cost
|
||||
|
||||
candidates = [model]
|
||||
if "/" in model:
|
||||
candidates.append(model.split("/", 1)[-1])
|
||||
|
||||
cost: Any = None
|
||||
for candidate in candidates:
|
||||
try:
|
||||
cost = completion_cost(
|
||||
completion_response={"model": candidate, "usage": usage_payload},
|
||||
model=model,
|
||||
)
|
||||
break
|
||||
except Exception: # nosec B112 # noqa: BLE001, S112
|
||||
continue
|
||||
|
||||
if cost is None:
|
||||
logger.debug("LiteLLM cost estimate unavailable for model %s", model)
|
||||
return None
|
||||
|
||||
return cost if isinstance(cost, int | float) and cost >= 0 else None
|
||||
|
||||
|
||||
def _litellm_model_name(model: str | None) -> str | None:
|
||||
if not model:
|
||||
return None
|
||||
normalized = model.strip()
|
||||
for prefix in ("litellm/", "any-llm/", "openai/"):
|
||||
if normalized.startswith(prefix):
|
||||
normalized = normalized.removeprefix(prefix)
|
||||
break
|
||||
return normalized or None
|
||||
|
||||
|
||||
def _details_to_dict(details: Any) -> dict[str, Any]:
|
||||
if details is None:
|
||||
return {}
|
||||
if isinstance(details, list):
|
||||
for item in details:
|
||||
result = _details_to_dict(item)
|
||||
if result:
|
||||
return result
|
||||
return {}
|
||||
if hasattr(details, "model_dump"):
|
||||
return _details_to_dict(details.model_dump())
|
||||
if not isinstance(details, dict):
|
||||
return {}
|
||||
return {str(k): v for k, v in details.items() if v is not None}
|
||||
|
||||
|
||||
def _int_or_zero(value: Any) -> int:
|
||||
try:
|
||||
return max(0, int(value or 0))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _float_or_zero(value: Any) -> float:
|
||||
try:
|
||||
result = float(value or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
return result if result >= 0 else 0.0
|
||||
|
||||
|
||||
def _round_cost(cost: float) -> float:
|
||||
return round(max(0.0, cost), 10)
|
||||
@@ -0,0 +1,197 @@
|
||||
"""Artifact writers for Strix scan reports."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import tempfile
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from strix.core.paths import run_record_path
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SEVERITY_ORDER = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
|
||||
|
||||
|
||||
def read_run_record(run_dir: Path) -> dict[str, Any]:
|
||||
path = run_record_path(run_dir)
|
||||
if not path.exists():
|
||||
return {}
|
||||
try:
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
raise RuntimeError(f"run.json at {path} is unreadable: {exc}") from exc
|
||||
if not isinstance(data, dict):
|
||||
raise TypeError(f"run.json at {path} is not an object")
|
||||
return data
|
||||
|
||||
|
||||
def write_run_record(run_dir: Path, run_record: dict[str, Any]) -> None:
|
||||
_atomic_write_text(
|
||||
run_record_path(run_dir),
|
||||
json.dumps(run_record, ensure_ascii=False, indent=2, default=str),
|
||||
)
|
||||
|
||||
|
||||
def write_executive_report(run_dir: Path, final_scan_result: str) -> None:
|
||||
path = run_dir / "penetration_test_report.md"
|
||||
with path.open("w", encoding="utf-8") as f:
|
||||
f.write("# Security Penetration Test Report\n\n")
|
||||
f.write(f"**Generated:** {datetime.now(UTC).strftime('%Y-%m-%d %H:%M:%S UTC')}\n\n")
|
||||
f.write(f"{final_scan_result}\n")
|
||||
logger.info("Saved final penetration test report to: %s", path)
|
||||
|
||||
|
||||
def write_vulnerabilities(
|
||||
run_dir: Path,
|
||||
vulnerability_reports: list[dict[str, Any]],
|
||||
saved_vuln_ids: set[str],
|
||||
) -> int:
|
||||
vuln_dir = run_dir / "vulnerabilities"
|
||||
vuln_dir.mkdir(exist_ok=True)
|
||||
|
||||
new_reports = [r for r in vulnerability_reports if r["id"] not in saved_vuln_ids]
|
||||
|
||||
for report in new_reports:
|
||||
_atomic_write_text(
|
||||
vuln_dir / f"{report['id']}.md",
|
||||
render_vulnerability_md(report),
|
||||
)
|
||||
saved_vuln_ids.add(report["id"])
|
||||
|
||||
sorted_reports = sorted(
|
||||
vulnerability_reports,
|
||||
key=lambda r: (_SEVERITY_ORDER.get(r["severity"], 5), r["timestamp"]),
|
||||
)
|
||||
csv_path = run_dir / "vulnerabilities.csv"
|
||||
csv_buf = io.StringIO()
|
||||
fieldnames = ["id", "title", "severity", "timestamp", "file"]
|
||||
csv_writer = csv.DictWriter(csv_buf, fieldnames=fieldnames, lineterminator="\r\n")
|
||||
csv_writer.writeheader()
|
||||
for report in sorted_reports:
|
||||
csv_writer.writerow(
|
||||
{
|
||||
"id": report["id"],
|
||||
"title": report["title"],
|
||||
"severity": report["severity"].upper(),
|
||||
"timestamp": report["timestamp"],
|
||||
"file": f"vulnerabilities/{report['id']}.md",
|
||||
},
|
||||
)
|
||||
_atomic_write_text(csv_path, csv_buf.getvalue())
|
||||
|
||||
_atomic_write_text(
|
||||
run_dir / "vulnerabilities.json",
|
||||
json.dumps(vulnerability_reports, ensure_ascii=False, indent=2, default=str),
|
||||
)
|
||||
|
||||
if new_reports:
|
||||
logger.info(
|
||||
"Saved %d new vulnerability report(s) to: %s",
|
||||
len(new_reports),
|
||||
vuln_dir,
|
||||
)
|
||||
logger.info("Updated vulnerability index: %s", csv_path)
|
||||
return len(new_reports)
|
||||
|
||||
|
||||
def _atomic_write_text(path: Path, payload: str) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
dir=str(path.parent),
|
||||
prefix=f".{path.name}.",
|
||||
suffix=".tmp",
|
||||
delete=False,
|
||||
) as tmp:
|
||||
tmp.write(payload)
|
||||
tmp_path = Path(tmp.name)
|
||||
tmp_path.replace(path)
|
||||
|
||||
|
||||
def render_vulnerability_md(report: dict[str, Any]) -> str: # noqa: PLR0912, PLR0915
|
||||
lines: list[str] = [
|
||||
f"# {report.get('title', 'Untitled Vulnerability')}\n",
|
||||
f"**ID:** {report.get('id', 'unknown')}",
|
||||
f"**Severity:** {report.get('severity', 'unknown').upper()}",
|
||||
f"**Found:** {report.get('timestamp', 'unknown')}",
|
||||
]
|
||||
|
||||
metadata: list[tuple[str, Any]] = [
|
||||
("Target", report.get("target")),
|
||||
("Endpoint", report.get("endpoint")),
|
||||
("Method", report.get("method")),
|
||||
("CVE", report.get("cve")),
|
||||
("CWE", report.get("cwe")),
|
||||
]
|
||||
cvss = report.get("cvss")
|
||||
if cvss is not None:
|
||||
metadata.append(("CVSS", cvss))
|
||||
for label, value in metadata:
|
||||
if value:
|
||||
lines.append(f"**{label}:** {value}")
|
||||
|
||||
lines.append("")
|
||||
lines.append("## Description\n")
|
||||
lines.append(report.get("description") or "No description provided.")
|
||||
lines.append("")
|
||||
|
||||
if report.get("impact"):
|
||||
lines.append("## Impact\n")
|
||||
lines.append(str(report["impact"]))
|
||||
lines.append("")
|
||||
|
||||
if report.get("technical_analysis"):
|
||||
lines.append("## Technical Analysis\n")
|
||||
lines.append(str(report["technical_analysis"]))
|
||||
lines.append("")
|
||||
|
||||
if report.get("poc_description") or report.get("poc_script_code"):
|
||||
lines.append("## Proof of Concept\n")
|
||||
if report.get("poc_description"):
|
||||
lines.append(str(report["poc_description"]))
|
||||
lines.append("")
|
||||
if report.get("poc_script_code"):
|
||||
lines.append("```")
|
||||
lines.append(str(report["poc_script_code"]))
|
||||
lines.append("```")
|
||||
lines.append("")
|
||||
|
||||
if report.get("code_locations"):
|
||||
lines.append("## Code Analysis\n")
|
||||
for i, loc in enumerate(report["code_locations"]):
|
||||
file_ref = loc.get("file", "unknown")
|
||||
line_ref = ""
|
||||
if loc.get("start_line") is not None:
|
||||
if loc.get("end_line") and loc["end_line"] != loc["start_line"]:
|
||||
line_ref = f" (lines {loc['start_line']}-{loc['end_line']})"
|
||||
else:
|
||||
line_ref = f" (line {loc['start_line']})"
|
||||
lines.append(f"**Location {i + 1}:** `{file_ref}`{line_ref}")
|
||||
if loc.get("label"):
|
||||
lines.append(f" {loc['label']}")
|
||||
if loc.get("snippet"):
|
||||
lines.append(f" ```\n {loc['snippet']}\n ```")
|
||||
if loc.get("fix_before") or loc.get("fix_after"):
|
||||
lines.append("\n **Suggested Fix:**")
|
||||
lines.append("```diff")
|
||||
if loc.get("fix_before"):
|
||||
lines.extend(f"- {ln}" for ln in str(loc["fix_before"]).splitlines())
|
||||
if loc.get("fix_after"):
|
||||
lines.extend(f"+ {ln}" for ln in str(loc["fix_after"]).splitlines())
|
||||
lines.append("```")
|
||||
lines.append("")
|
||||
|
||||
if report.get("remediation_steps"):
|
||||
lines.append("## Remediation\n")
|
||||
lines.append(str(report["remediation_steps"]))
|
||||
lines.append("")
|
||||
|
||||
return "\n".join(lines)
|
||||
Reference in New Issue
Block a user