fix(reports): auto-detect PoC language with Python fallback

This commit is contained in:
Ahmed Allam
2026-07-25 13:09:57 -07:00
committed by Ahmed Allam
parent 31c18f8f75
commit 97ed7e79a1
9 changed files with 211 additions and 165 deletions
+8 -8
View File
@@ -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("```")
+39
View File
@@ -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):