mirror of
https://github.com/usestrix/strix.git
synced 2026-08-16 09:26:39 +02:00
fix(reports): auto-detect PoC language with Python fallback
This commit is contained in:
@@ -39,7 +39,11 @@ 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.fenced import (
|
||||
guess_language_name,
|
||||
parse_fenced_code,
|
||||
resolve_lexer,
|
||||
)
|
||||
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
|
||||
@@ -333,14 +337,9 @@ class VulnerabilityDetailScreen(ModalScreen): # type: ignore[misc]
|
||||
|
||||
def _highlight_python(self, code: str, language: str | None = None) -> Text:
|
||||
try:
|
||||
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)
|
||||
lexer = resolve_lexer(language, code)
|
||||
style = get_style_by_name("native")
|
||||
colors = {
|
||||
token: f"#{style_def['color']}" for token, style_def in style if style_def["color"]
|
||||
@@ -608,7 +607,8 @@ class VulnerabilityDetailScreen(ModalScreen): # type: ignore[misc]
|
||||
lines.append("")
|
||||
if vuln.get("poc_script_code"):
|
||||
poc_language, poc_code = parse_fenced_code(vuln["poc_script_code"])
|
||||
lines.append(f"```{poc_language or 'python'}")
|
||||
fence_lang = poc_language or guess_language_name(poc_code)
|
||||
lines.append(f"```{fence_lang}")
|
||||
lines.append(poc_code)
|
||||
lines.append("```")
|
||||
|
||||
|
||||
@@ -1,5 +1,10 @@
|
||||
import re
|
||||
|
||||
from pygments.lexer import Lexer
|
||||
from pygments.lexers import PythonLexer, get_lexer_by_name, guess_lexer
|
||||
from pygments.lexers.special import TextLexer
|
||||
from pygments.util import ClassNotFound
|
||||
|
||||
|
||||
_FENCE_RE = re.compile(r"^```([^\n`]*)\n(.*?)\n?```$", re.DOTALL)
|
||||
|
||||
@@ -18,3 +23,37 @@ def parse_fenced_code(raw: str) -> tuple[str | None, str]:
|
||||
info = match.group(1).strip()
|
||||
language = info.split()[0] if info else None
|
||||
return (language or None), match.group(2)
|
||||
|
||||
|
||||
def resolve_lexer(language: str | None, code: str) -> Lexer:
|
||||
"""Pick a pygments lexer for ``code``.
|
||||
|
||||
Prefer the explicit fence ``language`` when it names a known lexer, otherwise
|
||||
auto-detect from the source. Fall back to Python when detection is
|
||||
inconclusive, since legacy (unfenced) PoC scripts are Python.
|
||||
"""
|
||||
if language:
|
||||
try:
|
||||
return get_lexer_by_name(language)
|
||||
except ClassNotFound:
|
||||
pass
|
||||
try:
|
||||
lexer = guess_lexer(code)
|
||||
except ClassNotFound:
|
||||
return PythonLexer()
|
||||
# ``guess_lexer`` returns the plain-text lexer when it can't detect anything.
|
||||
if isinstance(lexer, TextLexer):
|
||||
return PythonLexer()
|
||||
return lexer
|
||||
|
||||
|
||||
def guess_language_name(code: str) -> str:
|
||||
"""Return a markdown fence tag for ``code``, defaulting to ``python`` when
|
||||
auto-detection is inconclusive."""
|
||||
try:
|
||||
lexer = guess_lexer(code)
|
||||
except ClassNotFound:
|
||||
return "python"
|
||||
if isinstance(lexer, TextLexer) or not lexer.aliases:
|
||||
return "python"
|
||||
return str(lexer.aliases[0])
|
||||
|
||||
@@ -1,15 +1,12 @@
|
||||
from functools import cache
|
||||
from typing import Any, ClassVar
|
||||
|
||||
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 .fenced import parse_fenced_code, resolve_lexer
|
||||
from .registry import register_tool_renderer
|
||||
|
||||
|
||||
@@ -63,18 +60,9 @@ class CreateVulnerabilityReportRenderer(BaseToolRenderer):
|
||||
token_type = token_type.parent
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
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)
|
||||
lexer = resolve_lexer(language, code)
|
||||
text = Text()
|
||||
|
||||
for token_type, token_value in lexer.get_tokens(code):
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"use client";
|
||||
|
||||
import { useState } from "react";
|
||||
import hljs from "@/lib/hljs";
|
||||
import { highlightCode } from "@/lib/hljs";
|
||||
import "highlight.js/styles/github-dark.css";
|
||||
import { Copy, Check } from "lucide-react";
|
||||
import { copyToClipboard } from "@/lib/vulnerability-utils";
|
||||
@@ -35,16 +35,7 @@ export function MdCodeBlock({
|
||||
: fileName
|
||||
: null;
|
||||
|
||||
let highlighted: string;
|
||||
if (match) {
|
||||
try {
|
||||
highlighted = hljs.highlight(raw, { language: match[1], ignoreIllegals: true }).value;
|
||||
} catch {
|
||||
highlighted = hljs.highlightAuto(raw).value;
|
||||
}
|
||||
} else {
|
||||
highlighted = hljs.highlightAuto(raw).value;
|
||||
}
|
||||
const highlighted = highlightCode(raw, match?.[1]);
|
||||
|
||||
const lines = highlighted.split("\n");
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"use client";
|
||||
|
||||
import { useState } from "react";
|
||||
import hljs from "@/lib/hljs";
|
||||
import { highlightCode } from "@/lib/hljs";
|
||||
import "highlight.js/styles/github-dark.css";
|
||||
import ReactMarkdown from "react-markdown";
|
||||
import remarkGfm from "remark-gfm";
|
||||
@@ -22,16 +22,7 @@ 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 highlighted = highlightCode(code, language);
|
||||
|
||||
const copy = () => {
|
||||
if (!code) return;
|
||||
|
||||
@@ -11,4 +11,22 @@ hljs.registerLanguage("apache", apache);
|
||||
hljs.registerLanguage("dockerfile", dockerfile);
|
||||
hljs.registerLanguage("properties", properties);
|
||||
|
||||
/**
|
||||
* Highlight code, preferring an explicit language when it's recognized,
|
||||
* otherwise auto-detecting. Falls back to Python when auto-detection is
|
||||
* inconclusive, since legacy (unfenced) PoC scripts are Python.
|
||||
*/
|
||||
export function highlightCode(code: string, language?: string | null): string {
|
||||
try {
|
||||
if (language && hljs.getLanguage(language)) {
|
||||
return hljs.highlight(code, { language, ignoreIllegals: true }).value;
|
||||
}
|
||||
const auto = hljs.highlightAuto(code);
|
||||
if (auto.language) return auto.value;
|
||||
return hljs.highlight(code, { language: "python", ignoreIllegals: true }).value;
|
||||
} catch {
|
||||
return hljs.highlight(code, { language: "python", ignoreIllegals: true }).value;
|
||||
}
|
||||
}
|
||||
|
||||
export default hljs;
|
||||
|
||||
+119
-119
File diff suppressed because one or more lines are too long
@@ -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-B3PIcykV.js"></script>
|
||||
<script type="module" crossorigin src="./assets/index-B1jJGKIO.js"></script>
|
||||
<link rel="stylesheet" crossorigin href="./assets/index-vV8wxCG6.css">
|
||||
</head>
|
||||
<body>
|
||||
|
||||
@@ -2,7 +2,13 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from strix.interface.tui.renderers.fenced import parse_fenced_code
|
||||
from pygments.lexers import BashLexer, PythonLexer
|
||||
|
||||
from strix.interface.tui.renderers.fenced import (
|
||||
guess_language_name,
|
||||
parse_fenced_code,
|
||||
resolve_lexer,
|
||||
)
|
||||
from strix.viewer.report_pdf import _strip_code_fence
|
||||
|
||||
|
||||
@@ -43,3 +49,16 @@ def test_strip_code_fence_removes_fence() -> None:
|
||||
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"
|
||||
|
||||
|
||||
def test_resolve_lexer_honors_explicit_language() -> None:
|
||||
assert isinstance(resolve_lexer("bash", "echo hi"), BashLexer)
|
||||
|
||||
|
||||
def test_resolve_lexer_falls_back_to_python_when_unresolvable() -> None:
|
||||
# Unknown language name and empty body -> nothing to auto-detect -> Python.
|
||||
assert isinstance(resolve_lexer("not-a-language", ""), PythonLexer)
|
||||
|
||||
|
||||
def test_guess_language_name_defaults_to_python_when_inconclusive() -> None:
|
||||
assert guess_language_name("") == "python"
|
||||
|
||||
Reference in New Issue
Block a user