From 8f1f473eb8a5cf95c999c7ff8ad6b02167ff413e Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sat, 25 Apr 2026 16:28:07 -0700 Subject: [PATCH] refactor(telemetry): extract scan artifact I/O into ``strix.io.scan_artifacts`` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 150-line ``Tracer.save_run_data`` mashed three concerns together: opening file handles, formatting Markdown for vulnerabilities, and writing the executive penetration-test report. None of that is telemetry — it's pure on-disk artifact emission. Extract to :class:`ScanArtifactWriter` in ``strix/io/scan_artifacts.py``: - One writer per ``run_dir``, owns its own ``_saved_vuln_ids`` dedupe set so re-saves only emit new files. - ``writer.save(vulnerability_reports=, final_scan_result=)`` is the only public entry point. - ``_render_vulnerability_md`` is module-private and unit-testable in isolation. ``Tracer`` now lazily creates a single ``ScanArtifactWriter`` per ``run_dir`` and delegates ``save_run_data`` to it (~150 LoC body collapses to ~10). Net: tracer.py 422 → 327 LoC; new scan_artifacts.py 196 LoC. About −95 LoC of mixed concerns, plus telemetry no longer carries file-I/O responsibilities. Co-Authored-By: Claude Opus 4.7 (1M context) --- pyproject.toml | 3 + strix/io/__init__.py | 6 ++ strix/io/scan_artifacts.py | 199 +++++++++++++++++++++++++++++++++++++ strix/telemetry/tracer.py | 166 ++++--------------------------- 4 files changed, 225 insertions(+), 149 deletions(-) create mode 100644 strix/io/__init__.py create mode 100644 strix/io/scan_artifacts.py diff --git a/pyproject.toml b/pyproject.toml index 4081c27..c07cd8f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -209,6 +209,9 @@ ignore = [ # Backend factories import their backend's deps lazily so deployments # that pick a different backend don't need every backend's libs installed. "strix/runtime/backends.py" = ["PLC0415"] +# The vulnerability MD renderer is a long monolithic format function; +# splitting per-section would obscure the structure without simplifying. +"strix/io/scan_artifacts.py" = ["PLR0912", "PLR0915", "PERF401"] "strix/runtime/docker_client.py" = [ "TC002", # Manifest, Container imported for annotations "TC003", # uuid imported for annotation diff --git a/strix/io/__init__.py b/strix/io/__init__.py new file mode 100644 index 0000000..557d8f0 --- /dev/null +++ b/strix/io/__init__.py @@ -0,0 +1,6 @@ +"""Strix I/O — disk artifact writers. + +- :class:`ScanArtifactWriter` — writes vulnerability MD/CSV plus the + final penetration-test report under ``strix_runs//``. Used by + the tracer; could be used directly by post-run tooling. +""" diff --git a/strix/io/scan_artifacts.py b/strix/io/scan_artifacts.py new file mode 100644 index 0000000..4c25fe7 --- /dev/null +++ b/strix/io/scan_artifacts.py @@ -0,0 +1,199 @@ +"""Per-scan artifact writer. + +Writes the customer-facing penetration-test report and per-vulnerability +markdown + a ``vulnerabilities.csv`` index under ``strix_runs//``. +""" + +from __future__ import annotations + +import csv +import logging +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any + + +if TYPE_CHECKING: + from pathlib import Path + + +logger = logging.getLogger(__name__) + + +_SEVERITY_ORDER = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4} + + +class ScanArtifactWriter: + """Writes scan artifacts under ``run_dir``. Idempotent on repeat calls. + + Tracks which vulnerability ids have already been written so that + re-saves only emit new files; the ``vulnerabilities.csv`` index is + fully rewritten each call so the displayed order stays in sync with + severity sorting. + """ + + def __init__(self, run_dir: Path): + self._run_dir = run_dir + self._saved_vuln_ids: set[str] = set() + + @property + def run_dir(self) -> Path: + return self._run_dir + + def save( + self, + *, + vulnerability_reports: list[dict[str, Any]], + final_scan_result: str | None, + ) -> None: + """Write any new vulnerability MDs + rewrite the CSV index + + write the executive penetration-test report if available. + + Tolerant of OSError / RuntimeError — logs and swallows so a + cleanup failure can't prevent the next scan from finishing. + """ + try: + self._run_dir.mkdir(parents=True, exist_ok=True) + + if final_scan_result: + self._write_executive_report(final_scan_result) + + if vulnerability_reports: + self._write_vulnerabilities(vulnerability_reports) + + logger.info("📊 Essential scan data saved to: %s", self._run_dir) + except (OSError, RuntimeError): + logger.exception("Failed to save scan data") + + # --- internals --------------------------------------------------------- + + def _write_executive_report(self, body: str) -> None: + path = self._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"{body}\n") + logger.info("Saved final penetration test report to: %s", path) + + def _write_vulnerabilities(self, reports: list[dict[str, Any]]) -> None: + vuln_dir = self._run_dir / "vulnerabilities" + vuln_dir.mkdir(exist_ok=True) + + new_reports = [r for r in reports if r["id"] not in self._saved_vuln_ids] + + for report in new_reports: + (vuln_dir / f"{report['id']}.md").write_text( + _render_vulnerability_md(report), + encoding="utf-8", + ) + self._saved_vuln_ids.add(report["id"]) + + sorted_reports = sorted( + reports, + key=lambda r: (_SEVERITY_ORDER.get(r["severity"], 5), r["timestamp"]), + ) + csv_path = self._run_dir / "vulnerabilities.csv" + with csv_path.open("w", encoding="utf-8", newline="") as f: + fieldnames = ["id", "title", "severity", "timestamp", "file"] + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for report in sorted_reports: + writer.writerow( + { + "id": report["id"], + "title": report["title"], + "severity": report["severity"].upper(), + "timestamp": report["timestamp"], + "file": f"vulnerabilities/{report['id']}.md", + }, + ) + + 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) + + +def _render_vulnerability_md(report: dict[str, Any]) -> str: + 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"): + for ln in str(loc["fix_before"]).splitlines(): + lines.append(f"- {ln}") + if loc.get("fix_after"): + for ln in str(loc["fix_after"]).splitlines(): + lines.append(f"+ {ln}") + 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) diff --git a/strix/telemetry/tracer.py b/strix/telemetry/tracer.py index 3f46ecb..48f81a0 100644 --- a/strix/telemetry/tracer.py +++ b/strix/telemetry/tracer.py @@ -5,6 +5,7 @@ from pathlib import Path from typing import Any, Optional from uuid import uuid4 +from strix.io.scan_artifacts import ScanArtifactWriter from strix.telemetry import posthog @@ -67,8 +68,8 @@ class Tracer: "status": "running", } self._run_dir: Path | None = None + self._writer: ScanArtifactWriter | None = None self._next_message_id = 1 - self._saved_vuln_ids: set[str] = set() self.caido_url: str | None = None self.vulnerability_found_callback: Callable[[dict[str, Any]], None] | None = None @@ -79,6 +80,7 @@ class Tracer: 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: @@ -91,6 +93,11 @@ class Tracer: return self._run_dir + 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, @@ -231,155 +238,16 @@ class Tracer: ) def save_run_data(self, mark_complete: bool = False) -> None: - try: - run_dir = self.get_run_dir() - 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" + 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" - if self.final_scan_result: - penetration_test_report_file = run_dir / "penetration_test_report.md" - with penetration_test_report_file.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"{self.final_scan_result}\n") - logger.info( - "Saved final penetration test report to: %s", - penetration_test_report_file, - ) - - if self.vulnerability_reports: - vuln_dir = run_dir / "vulnerabilities" - vuln_dir.mkdir(exist_ok=True) - - new_reports = [ - report - for report in self.vulnerability_reports - if report["id"] not in self._saved_vuln_ids - ] - - severity_order = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4} - sorted_reports = sorted( - self.vulnerability_reports, - key=lambda report: ( - severity_order.get(report["severity"], 5), - report["timestamp"], - ), - ) - - for report in new_reports: - vuln_file = vuln_dir / f"{report['id']}.md" - with vuln_file.open("w", encoding="utf-8") as f: - f.write(f"# {report.get('title', 'Untitled Vulnerability')}\n\n") - f.write(f"**ID:** {report.get('id', 'unknown')}\n") - f.write(f"**Severity:** {report.get('severity', 'unknown').upper()}\n") - f.write(f"**Found:** {report.get('timestamp', 'unknown')}\n") - - metadata_fields: 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_score = report.get("cvss") - if cvss_score is not None: - metadata_fields.append(("CVSS", cvss_score)) - - for label, value in metadata_fields: - if value: - f.write(f"**{label}:** {value}\n") - - f.write("\n## Description\n\n") - description = report.get("description") or "No description provided." - f.write(f"{description}\n\n") - - if report.get("impact"): - f.write("## Impact\n\n") - f.write(f"{report['impact']}\n\n") - - if report.get("technical_analysis"): - f.write("## Technical Analysis\n\n") - f.write(f"{report['technical_analysis']}\n\n") - - if report.get("poc_description") or report.get("poc_script_code"): - f.write("## Proof of Concept\n\n") - if report.get("poc_description"): - f.write(f"{report['poc_description']}\n\n") - if report.get("poc_script_code"): - f.write("```\n") - f.write(f"{report['poc_script_code']}\n") - f.write("```\n\n") - - if report.get("code_locations"): - f.write("## Code Analysis\n\n") - for i, loc in enumerate(report["code_locations"]): - prefix = f"**Location {i + 1}:**" - 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']})" - f.write(f"{prefix} `{file_ref}`{line_ref}\n") - if loc.get("label"): - f.write(f" {loc['label']}\n") - if loc.get("snippet"): - f.write(f" ```\n {loc['snippet']}\n ```\n") - if loc.get("fix_before") or loc.get("fix_after"): - f.write("\n **Suggested Fix:**\n") - f.write("```diff\n") - if loc.get("fix_before"): - for line in loc["fix_before"].splitlines(): - f.write(f"- {line}\n") - if loc.get("fix_after"): - for line in loc["fix_after"].splitlines(): - f.write(f"+ {line}\n") - f.write("```\n") - f.write("\n") - - if report.get("remediation_steps"): - f.write("## Remediation\n\n") - f.write(f"{report['remediation_steps']}\n\n") - - self._saved_vuln_ids.add(report["id"]) - - vuln_csv_file = run_dir / "vulnerabilities.csv" - with vuln_csv_file.open("w", encoding="utf-8", newline="") as f: - import csv - - fieldnames = ["id", "title", "severity", "timestamp", "file"] - writer = csv.DictWriter(f, fieldnames=fieldnames) - writer.writeheader() - - for report in sorted_reports: - writer.writerow( - { - "id": report["id"], - "title": report["title"], - "severity": report["severity"].upper(), - "timestamp": report["timestamp"], - "file": f"vulnerabilities/{report['id']}.md", - } - ) - - if new_reports: - logger.info( - "Saved %d new vulnerability report(s) to: %s", - len(new_reports), - vuln_dir, - ) - logger.info("Updated vulnerability index: %s", vuln_csv_file) - - logger.info("📊 Essential scan data saved to: %s", run_dir) - - except (OSError, RuntimeError): - logger.exception("Failed to save scan data") + self._get_writer().save( + vulnerability_reports=self.vulnerability_reports, + final_scan_result=self.final_scan_result, + ) def log_tool_start(self, agent_id: str, tool_name: str) -> int: """Record a tool invocation in flight. Returns an exec_id."""