fix(reports): strip markdown code fence from poc_script_code before rendering

2026-08-16 09:26:39 +02:00 · 2026-07-25 13:09:57 -07:00
parent f23fadfbff
commit 31c18f8f75
11 changed files with 334 additions and 153 deletions
@@ -122,6 +122,7 @@ module = [
    "pydantic_settings.*",
    "reportlab.*",
    "pypdf.*",
+    "pygments.*",
 ]
 ignore_missing_imports = true
 disable_error_code = ["import-untyped"]
@@ -39,6 +39,7 @@ from strix.interface.tui.live_view import TuiLiveView
 from strix.interface.tui.messages import send_user_message_to_agent
 from strix.interface.tui.renderers import render_tool_widget
 from strix.interface.tui.renderers.agent_message_renderer import AgentMessageRenderer
+from strix.interface.tui.renderers.fenced import parse_fenced_code
 from strix.interface.tui.renderers.user_message_renderer import UserMessageRenderer
 from strix.interface.utils import build_tui_stats_text
 from strix.report.state import ReportState, set_global_report_state
@@ -330,12 +331,16 @@ class VulnerabilityDetailScreen(ModalScreen):  # type: ignore[misc]
            return "#65a30d"
        return "#6b7280"

-    def _highlight_python(self, code: str) -> Text:
+    def _highlight_python(self, code: str, language: str | None = None) -> Text:
        try:
-            from pygments.lexers import PythonLexer
+            from pygments.lexers import PythonLexer, get_lexer_by_name
            from pygments.styles import get_style_by_name
+            from pygments.util import ClassNotFound

            lexer = PythonLexer()
+            if language:
+                with contextlib.suppress(ClassNotFound):
+                    lexer = get_lexer_by_name(language)
            style = get_style_by_name("native")
            colors = {
                token: f"#{style_def['color']}" for token, style_def in style if style_def["color"]
@@ -501,10 +506,11 @@ class VulnerabilityDetailScreen(ModalScreen):  # type: ignore[misc]

        poc_script_code = vuln.get("poc_script_code", "")
        if poc_script_code:
+            poc_language, poc_code = parse_fenced_code(poc_script_code)
            text.append("\n\n")
            text.append("PoC Code", style=self.FIELD_STYLE)
            text.append("\n")
-            text.append_text(self._highlight_python(poc_script_code))
+            text.append_text(self._highlight_python(poc_code, poc_language))

        remediation_steps = vuln.get("remediation_steps", "")
        if remediation_steps:
@@ -601,8 +607,9 @@ class VulnerabilityDetailScreen(ModalScreen):  # type: ignore[misc]
                lines.append(vuln["poc_description"])
                lines.append("")
            if vuln.get("poc_script_code"):
-                lines.append("```python")
-                lines.append(vuln["poc_script_code"])
+                poc_language, poc_code = parse_fenced_code(vuln["poc_script_code"])
+                lines.append(f"```{poc_language or 'python'}")
+                lines.append(poc_code)
                lines.append("```")

        if vuln.get("code_locations"):
@@ -0,0 +1,20 @@
+import re
+
+
+_FENCE_RE = re.compile(r"^```([^\n`]*)\n(.*?)\n?```$", re.DOTALL)
+
+
+def parse_fenced_code(raw: str) -> tuple[str | None, str]:
+    """Split an optionally fenced code string into ``(language, code)``.
+
+    Agent-generated code fields (e.g. ``poc_script_code``) are stored wrapped in
+    a markdown fence carrying the language, like ``` ```python\n...\n``` ```.
+    Return the fence's language tag and the inner code, or ``(None, raw)`` when
+    the value isn't fenced.
+    """
+    match = _FENCE_RE.match(raw.strip())
+    if not match:
+        return None, raw
+    info = match.group(1).strip()
+    language = info.split()[0] if info else None
+    return (language or None), match.group(2)
@@ -1,12 +1,15 @@
 from functools import cache
 from typing import Any, ClassVar

-from pygments.lexers import PythonLexer
+from pygments.lexer import Lexer
+from pygments.lexers import PythonLexer, get_lexer_by_name
 from pygments.styles import get_style_by_name
+from pygments.util import ClassNotFound
 from rich.text import Text
 from textual.widgets import Static

 from .base_renderer import BaseToolRenderer
+from .fenced import parse_fenced_code
 from .registry import register_tool_renderer


@@ -61,8 +64,17 @@ class CreateVulnerabilityReportRenderer(BaseToolRenderer):
        return None

    @classmethod
-    def _highlight_python(cls, code: str) -> Text:
-        lexer = PythonLexer()
+    def _get_lexer(cls, language: str | None) -> Lexer:
+        if language:
+            try:
+                return get_lexer_by_name(language)
+            except ClassNotFound:
+                pass
+        return PythonLexer()
+
+    @classmethod
+    def _highlight_code(cls, code: str, language: str | None) -> Text:
+        lexer = cls._get_lexer(language)
        text = Text()

        for token_type, token_value in lexer.get_tokens(code):
@@ -234,10 +246,11 @@ class CreateVulnerabilityReportRenderer(BaseToolRenderer):
            text.append(poc_description)

        if poc_script_code:
+            poc_language, poc_code = parse_fenced_code(poc_script_code)
            text.append("\n\n")
            text.append("PoC Code", style=FIELD_STYLE)
            text.append("\n")
-            text.append_text(cls._highlight_python(poc_script_code))
+            text.append_text(cls._highlight_code(poc_code, poc_language))

        if remediation_steps:
            text.append("\n\n")
@@ -3,6 +3,7 @@
 import type { ToolRendererProps } from "@/types/events";
 import { TruncatedText } from "./ToolCard";
 import { MdCodeBlock } from "@/components/vulnerability/MdCodeBlock";
+import { parseFencedCode } from "@/lib/fenced-code";
 import Markdown from "./Markdown";

 const SEVERITY_COLORS: Record<string, string> = {
@@ -19,7 +20,7 @@ export default function VulnReportRenderer({ args, result }: ToolRendererProps)
  const method = (args.method as string) ?? "";
  const technicalAnalysis = (args.technical_analysis as string) ?? "";
  const pocDescription = (args.poc_description as string) ?? "";
-  const pocCode = (args.poc_script_code as string) ?? "";
+  const { language: pocLang, code: pocCode } = parseFencedCode((args.poc_script_code as string) ?? "");
  const remediation = (args.remediation_steps as string) ?? "";
  const cve = (args.cve as string) ?? "";
  const cwe = (args.cwe as string) ?? "";
@@ -59,7 +60,7 @@ export default function VulnReportRenderer({ args, result }: ToolRendererProps)
        <div>
          <span className="text-emerald-400/60 text-sm font-semibold">Proof of Concept</span>
          {pocDescription && <div className="mt-1"><Markdown text={pocDescription} /></div>}
-          {pocCode && <MdCodeBlock>{pocCode}</MdCodeBlock>}
+          {pocCode && <MdCodeBlock className={pocLang ? `language-${pocLang}` : undefined}>{pocCode}</MdCodeBlock>}
        </div>
      )}
      {remediation && (
@@ -7,6 +7,7 @@ import ReactMarkdown from "react-markdown";
 import remarkGfm from "remark-gfm";
 import { Copy, Check } from "lucide-react";
 import { copyToClipboard } from "@/lib/vulnerability-utils";
+import { parseFencedCode } from "@/lib/fenced-code";
 import { rehypeCodeMeta, mdComponents } from "./MdCodeBlock";

 interface PocBlockProps {
@@ -20,9 +21,21 @@ export function PocBlock({ description, scriptCode, onCopy }: PocBlockProps) {

  if (!description && !scriptCode) return null;

+  const { language, code } = parseFencedCode(scriptCode);
+  const highlighted = (() => {
+    try {
+      if (language && hljs.getLanguage(language)) {
+        return hljs.highlight(code, { language }).value;
+      }
+      return hljs.highlightAuto(code).value;
+    } catch {
+      return hljs.highlightAuto(code).value;
+    }
+  })();
+
  const copy = () => {
-    if (!scriptCode) return;
-    copyToClipboard(scriptCode);
+    if (!code) return;
+    copyToClipboard(code);
    setCopied(true);
    setTimeout(() => setCopied(false), 2000);
    onCopy?.();
@@ -43,7 +56,7 @@ export function PocBlock({ description, scriptCode, onCopy }: PocBlockProps) {
            </ReactMarkdown>
          </div>
        )}
-        {scriptCode && (
+        {code && (
          <div className="group/poc relative rounded-md border border-[#2a2a2a] overflow-hidden">
            <div className="flex items-stretch">
              <span className="relative flex items-center text-[13px] text-[#999] font-mono px-4 py-2 border-r border-[#2a2a2a]">PoC Script<span className="absolute top-0 inset-x-0 h-0.5 bg-white/60 rounded-full" /></span>
@@ -64,7 +77,7 @@ export function PocBlock({ description, scriptCode, onCopy }: PocBlockProps) {
              <pre className="font-mono text-[12px] leading-[22px] whitespace-pre [font-variant-ligatures:none]">
                <code
                  dangerouslySetInnerHTML={{
-                    __html: hljs.highlight(scriptCode, { language: "python" }).value,
+                    __html: highlighted,
                  }}
                />
              </pre>
@@ -0,0 +1,27 @@
+export interface ParsedFencedCode {
+  language?: string;
+  code: string;
+}
+
+const FENCE_RE = /^```([^\n`]*)\n([\s\S]*?)\n?```$/;
+
+/**
+ * Agent-generated `poc_script_code` is stored wrapped in a markdown code fence
+ * that carries the language, e.g.
+ *
+ *     ```python
+ *     import requests
+ *     ```
+ *
+ * Renderers that show the value as bare code must not display the fence lines
+ * literally. This extracts the inner code and the fence's language tag. Returns
+ * the input unchanged (no language) when it isn't fenced.
+ */
+export function parseFencedCode(raw: string | null | undefined): ParsedFencedCode {
+  if (!raw) return { code: "" };
+  const match = FENCE_RE.exec(raw.trim());
+  if (!match) return { code: raw };
+  const info = match[1].trim();
+  const language = info ? info.split(/\s+/)[0] : undefined;
+  return { language: language || undefined, code: match[2] };
+}
@@ -151,19 +151,35 @@ def _styles() -> dict[str, ParagraphStyle]:
        "Finding", fontName=_SANS_BOLD, fontSize=13, leading=17, textColor=_INK, spaceBefore=6
    )
    styles["field_label"] = ParagraphStyle(
-        "FieldLabel", fontName=_SANS_BOLD, fontSize=8.5, leading=12, textColor=_MUTED,
-        spaceBefore=10, spaceAfter=2,
+        "FieldLabel",
+        fontName=_SANS_BOLD,
+        fontSize=8.5,
+        leading=12,
+        textColor=_MUTED,
+        spaceBefore=10,
+        spaceAfter=2,
    )
    styles["body"] = ParagraphStyle(
        "Body", fontName=_SANS, fontSize=10, leading=15, textColor=_TEXT, spaceAfter=8
    )
    styles["md_heading"] = ParagraphStyle(
-        "MdHeading", fontName=_SANS_BOLD, fontSize=11, leading=15, textColor=_INK,
-        spaceBefore=10, spaceAfter=4,
+        "MdHeading",
+        fontName=_SANS_BOLD,
+        fontSize=11,
+        leading=15,
+        textColor=_INK,
+        spaceBefore=10,
+        spaceAfter=4,
    )
    styles["bullet"] = ParagraphStyle(
-        "Bullet", fontName=_SANS, fontSize=10, leading=15, textColor=_TEXT,
-        leftIndent=16, firstLineIndent=-11, spaceAfter=3,
+        "Bullet",
+        fontName=_SANS,
+        fontSize=10,
+        leading=15,
+        textColor=_TEXT,
+        leftIndent=16,
+        firstLineIndent=-11,
+        spaceAfter=3,
    )
    styles["meta_inline"] = ParagraphStyle(
        "MetaInline", fontName=_SANS, fontSize=9, leading=13, textColor=_MUTED, spaceBefore=4
@@ -172,23 +188,44 @@ def _styles() -> dict[str, ParagraphStyle]:
    # a bordered paragraph's top padding, so too small a gap lets the background
    # box bleed up over the field label above it.
    styles["code"] = ParagraphStyle(
-        "Code", fontName=_MONO, fontSize=8, leading=11, textColor=_TEXT,
-        backColor=_LIGHT_BG, borderColor=_BORDER, borderWidth=0.5, borderPadding=8,
-        spaceBefore=12, spaceAfter=12,
+        "Code",
+        fontName=_MONO,
+        fontSize=8,
+        leading=11,
+        textColor=_TEXT,
+        backColor=_LIGHT_BG,
+        borderColor=_BORDER,
+        borderWidth=0.5,
+        borderPadding=8,
+        spaceBefore=12,
+        spaceAfter=12,
    )
    styles["count"] = ParagraphStyle(
        "Count", fontName=_SANS_BOLD, fontSize=30, leading=32, alignment=TA_CENTER
    )
    styles["count_label"] = ParagraphStyle(
-        "CountLabel", fontName=_SANS_BOLD, fontSize=8, leading=12, textColor=_MUTED,
-        alignment=TA_CENTER, spaceBefore=4,
+        "CountLabel",
+        fontName=_SANS_BOLD,
+        fontSize=8,
+        leading=12,
+        textColor=_MUTED,
+        alignment=TA_CENTER,
+        spaceBefore=4,
    )
    styles["badge"] = ParagraphStyle(
-        "Badge", fontName=_SANS_BOLD, fontSize=9, leading=11, textColor=colors.white,
+        "Badge",
+        fontName=_SANS_BOLD,
+        fontSize=9,
+        leading=11,
+        textColor=colors.white,
        alignment=TA_CENTER,
    )
    styles["confidential"] = ParagraphStyle(
-        "Confidential", fontName=_SANS_BOLD, fontSize=9, leading=12, textColor=colors.white,
+        "Confidential",
+        fontName=_SANS_BOLD,
+        fontSize=9,
+        leading=12,
+        textColor=colors.white,
        alignment=TA_CENTER,
    )
    return styles
@@ -256,8 +293,10 @@ def _severity_grid(styles: dict[str, ParagraphStyle], counts: dict[str, int]) ->
        color = _SEVERITY_COLORS[name]
        count_style = ParagraphStyle(f"Count{name}", parent=styles["count"], textColor=color)
        cells.append(
-            [Paragraph(str(counts.get(name, 0)), count_style),
-             Paragraph(name.upper(), styles["count_label"])]
+            [
+                Paragraph(str(counts.get(name, 0)), count_style),
+                Paragraph(name.upper(), styles["count_label"]),
+            ]
        )
    col = (_PAGE_W - 40 * mm) / 4
    table = Table([cells], colWidths=[col] * 4)
@@ -320,8 +359,10 @@ def _cover(
        ("DURATION", _duration(record.get("start_time"), record.get("end_time"))),
    ]
    meta_table = Table(
-        [[Paragraph(label, styles["meta_label"]), Paragraph(_esc(value), styles["meta_value"])]
-         for label, value in meta_rows],
+        [
+            [Paragraph(label, styles["meta_label"]), Paragraph(_esc(value), styles["meta_value"])]
+            for label, value in meta_rows
+        ],
        colWidths=[38 * mm, _PAGE_W - 40 * mm - 38 * mm],
    )
    meta_table.setStyle(
@@ -462,6 +503,18 @@ def _markdown_flowables(  # noqa: PLR0915 - cohesive block parser, splitting hur
    return flow


+_FENCE_RE = re.compile(r"^```([^\n`]*)\n(.*?)\n?```$", re.DOTALL)
+
+
+def _strip_code_fence(value: Any) -> Any:
+    """Drop a wrapping markdown code fence so raw-code fields don't show the
+    ``` ```lang ``` marker lines literally in the PDF."""
+    if not isinstance(value, str):
+        return value
+    match = _FENCE_RE.match(value.strip())
+    return match.group(2) if match else value
+
+
 def _field_block(
    styles: dict[str, ParagraphStyle], label: str, value: Any, *, code: bool = False
 ) -> list[Flowable]:
@@ -503,7 +556,8 @@ def _finding_flowables(
    story.extend(_field_block(styles, "Impact", vuln.get("impact")))
    story.extend(_field_block(styles, "Technical analysis", vuln.get("technical_analysis")))
    story.extend(_field_block(styles, "Proof of concept", vuln.get("poc_description")))
-    story.extend(_field_block(styles, "PoC script", vuln.get("poc_script_code"), code=True))
+    poc_script = _strip_code_fence(vuln.get("poc_script_code"))
+    story.extend(_field_block(styles, "PoC script", poc_script, code=True))
    story.extend(_field_block(styles, "Evidence", vuln.get("evidence"), code=True))

    remediation = vuln.get("remediation_steps")
@@ -6,7 +6,7 @@
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <meta name="color-scheme" content="dark" />
    <title>Strix Results</title>
-    <script type="module" crossorigin src="./assets/index-e2r6VuTm.js"></script>
+    <script type="module" crossorigin src="./assets/index-B3PIcykV.js"></script>
    <link rel="stylesheet" crossorigin href="./assets/index-vV8wxCG6.css">
  </head>
  <body>
@@ -0,0 +1,45 @@
+"""Tests for stripping the markdown code fence off stored code fields."""
+
+from __future__ import annotations
+
+from strix.interface.tui.renderers.fenced import parse_fenced_code
+from strix.viewer.report_pdf import _strip_code_fence
+
+
+def test_parse_fenced_code_extracts_language_and_body() -> None:
+    language, code = parse_fenced_code("```python\nimport requests\nprint(1)\n```")
+    assert language == "python"
+    assert code == "import requests\nprint(1)"
+
+
+def test_parse_fenced_code_uses_first_token_of_info_string() -> None:
+    language, code = parse_fenced_code("```python title=app.py\nx = 1\n```")
+    assert language == "python"
+    assert code == "x = 1"
+
+
+def test_parse_fenced_code_handles_non_python_language() -> None:
+    language, code = parse_fenced_code("```http\nGET / HTTP/1.1\n```")
+    assert language == "http"
+    assert code == "GET / HTTP/1.1"
+
+
+def test_parse_fenced_code_passes_through_unfenced() -> None:
+    language, code = parse_fenced_code("import requests\nprint(1)")
+    assert language is None
+    assert code == "import requests\nprint(1)"
+
+
+def test_parse_fenced_code_fence_without_language() -> None:
+    language, code = parse_fenced_code("```\nplain\n```")
+    assert language is None
+    assert code == "plain"
+
+
+def test_strip_code_fence_removes_fence() -> None:
+    assert _strip_code_fence("```python\nx = 1\n```") == "x = 1"
+
+
+def test_strip_code_fence_passes_through_non_string_and_unfenced() -> None:
+    assert _strip_code_fence(None) is None
+    assert _strip_code_fence("x = 1") == "x = 1"