fix: markdown rendering, continuous polling, and attack chain graph loading
- Add markdown-to-HTML renderer for assistant messages (headers, bold, code blocks, lists, inline code) - Fix polling to continuously loop while session is running using poll_gen signal - Fix attack chain graph loading with spawn delay for DOM readiness - Default attack chain tab to list view (more reliable) - Render tool_result role messages as tool indicators Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -8,6 +8,157 @@ use crate::infrastructure::pentest::{
|
||||
fetch_pentest_session, send_pentest_message,
|
||||
};
|
||||
|
||||
/// Simple markdown-to-HTML converter for assistant messages.
|
||||
/// Handles headers, bold, italic, code blocks, inline code, and lists.
|
||||
fn markdown_to_html(input: &str) -> String {
|
||||
let mut html = String::new();
|
||||
let mut in_code_block = false;
|
||||
let mut in_list = false;
|
||||
|
||||
for line in input.lines() {
|
||||
if line.starts_with("```") {
|
||||
if in_code_block {
|
||||
html.push_str("</code></pre>");
|
||||
in_code_block = false;
|
||||
} else {
|
||||
if in_list {
|
||||
html.push_str("</ul>");
|
||||
in_list = false;
|
||||
}
|
||||
html.push_str("<pre style=\"background:var(--bg-primary);padding:10px;border-radius:6px;overflow-x:auto;font-size:0.8rem;margin:6px 0;\"><code>");
|
||||
in_code_block = true;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if in_code_block {
|
||||
// Escape HTML inside code blocks
|
||||
let escaped = line
|
||||
.replace('&', "&")
|
||||
.replace('<', "<")
|
||||
.replace('>', ">");
|
||||
html.push_str(&escaped);
|
||||
html.push('\n');
|
||||
continue;
|
||||
}
|
||||
|
||||
let trimmed = line.trim();
|
||||
|
||||
// Blank line — close list if open
|
||||
if trimmed.is_empty() {
|
||||
if in_list {
|
||||
html.push_str("</ul>");
|
||||
in_list = false;
|
||||
}
|
||||
html.push_str("<br/>");
|
||||
continue;
|
||||
}
|
||||
|
||||
// Lists
|
||||
if trimmed.starts_with("- ") || trimmed.starts_with("* ") {
|
||||
if !in_list {
|
||||
html.push_str("<ul style=\"margin:4px 0;padding-left:20px;\">");
|
||||
in_list = true;
|
||||
}
|
||||
let content = inline_format(&trimmed[2..]);
|
||||
html.push_str(&format!("<li>{content}</li>"));
|
||||
continue;
|
||||
}
|
||||
|
||||
// Numbered lists
|
||||
if trimmed.len() > 2 {
|
||||
let mut chars = trimmed.chars();
|
||||
let first = chars.next();
|
||||
let second = chars.next();
|
||||
if first.map(|c| c.is_ascii_digit()).unwrap_or(false)
|
||||
&& (second == Some('.') || second == Some(')'))
|
||||
{
|
||||
let rest = &trimmed[2..].trim_start();
|
||||
if !in_list {
|
||||
html.push_str("<ul style=\"margin:4px 0;padding-left:20px;\">");
|
||||
in_list = true;
|
||||
}
|
||||
let content = inline_format(rest);
|
||||
html.push_str(&format!("<li>{content}</li>"));
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Close list if we're no longer in one
|
||||
if in_list {
|
||||
html.push_str("</ul>");
|
||||
in_list = false;
|
||||
}
|
||||
|
||||
// Headers
|
||||
if trimmed.starts_with("### ") {
|
||||
let content = inline_format(&trimmed[4..]);
|
||||
html.push_str(&format!(
|
||||
"<h4 style=\"margin:8px 0 4px;font-size:0.9rem;\">{content}</h4>"
|
||||
));
|
||||
} else if trimmed.starts_with("## ") {
|
||||
let content = inline_format(&trimmed[3..]);
|
||||
html.push_str(&format!(
|
||||
"<h3 style=\"margin:10px 0 4px;font-size:0.95rem;\">{content}</h3>"
|
||||
));
|
||||
} else if trimmed.starts_with("# ") {
|
||||
let content = inline_format(&trimmed[2..]);
|
||||
html.push_str(&format!(
|
||||
"<h2 style=\"margin:12px 0 6px;font-size:1rem;\">{content}</h2>"
|
||||
));
|
||||
} else {
|
||||
let content = inline_format(trimmed);
|
||||
html.push_str(&format!("<p style=\"margin:2px 0;\">{content}</p>"));
|
||||
}
|
||||
}
|
||||
|
||||
if in_list {
|
||||
html.push_str("</ul>");
|
||||
}
|
||||
if in_code_block {
|
||||
html.push_str("</code></pre>");
|
||||
}
|
||||
|
||||
html
|
||||
}
|
||||
|
||||
/// Handle inline formatting: bold, italic, inline code
|
||||
fn inline_format(text: &str) -> String {
|
||||
let mut result = text
|
||||
.replace('&', "&")
|
||||
.replace('<', "<")
|
||||
.replace('>', ">");
|
||||
|
||||
// Inline code (backticks)
|
||||
while let Some(start) = result.find('`') {
|
||||
if let Some(end) = result[start + 1..].find('`') {
|
||||
let code_content = &result[start + 1..start + 1 + end].to_string();
|
||||
let replacement = format!(
|
||||
"<code style=\"background:var(--bg-primary);padding:1px 4px;border-radius:3px;font-size:0.85em;\">{code_content}</code>"
|
||||
);
|
||||
result = format!("{}{}{}", &result[..start], replacement, &result[start + 2 + end..]);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Bold (**text**)
|
||||
while let Some(start) = result.find("**") {
|
||||
if let Some(end) = result[start + 2..].find("**") {
|
||||
let bold_content = &result[start + 2..start + 2 + end].to_string();
|
||||
result = format!(
|
||||
"{}<strong>{bold_content}</strong>{}",
|
||||
&result[..start],
|
||||
&result[start + 4 + end..]
|
||||
);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
#[component]
|
||||
pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
let sid = session_id.clone();
|
||||
@@ -35,10 +186,11 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
let mut input_text = use_signal(String::new);
|
||||
let mut sending = use_signal(|| false);
|
||||
let mut right_tab = use_signal(|| "findings".to_string());
|
||||
let mut chain_view = use_signal(|| "graph".to_string()); // "graph" or "list"
|
||||
let mut chain_view = use_signal(|| "list".to_string());
|
||||
let mut exporting = use_signal(|| false);
|
||||
let mut poll_gen = use_signal(|| 0u32); // incremented to trigger re-poll
|
||||
|
||||
// Auto-poll messages every 3s when session is running
|
||||
// Extract session status
|
||||
let session_status = {
|
||||
let s = session.read();
|
||||
match &*s {
|
||||
@@ -54,10 +206,10 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
|
||||
let is_running = session_status == "running";
|
||||
|
||||
let sid_for_poll = session_id.clone();
|
||||
// Continuous polling: re-fetch all data every 3s while running
|
||||
use_effect(move || {
|
||||
let _gen = *poll_gen.read(); // subscribe to changes
|
||||
if is_running {
|
||||
let _sid = sid_for_poll.clone();
|
||||
spawn(async move {
|
||||
#[cfg(feature = "web")]
|
||||
gloo_timers::future::TimeoutFuture::new(3_000).await;
|
||||
@@ -67,24 +219,36 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
findings.restart();
|
||||
attack_chain.restart();
|
||||
session.restart();
|
||||
// Bump generation to trigger the next poll cycle
|
||||
let next = poll_gen.peek().wrapping_add(1);
|
||||
poll_gen.set(next);
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// Load attack chain into vis-network when data changes and graph tab is active
|
||||
let chain_data_for_viz = attack_chain.read().clone();
|
||||
let current_tab = right_tab.read().clone();
|
||||
let current_chain_view = chain_view.read().clone();
|
||||
// Load attack chain into vis-network when graph tab is active and data is available
|
||||
// Use a separate effect that reads the reactive resources directly
|
||||
use_effect(move || {
|
||||
if current_tab == "chain" && current_chain_view == "graph" {
|
||||
if let Some(Some(data)) = &chain_data_for_viz {
|
||||
let tab = right_tab.read().clone();
|
||||
let view = chain_view.read().clone();
|
||||
let chain = attack_chain.read().clone();
|
||||
|
||||
if tab == "chain" && view == "graph" {
|
||||
if let Some(Some(data)) = &chain {
|
||||
if !data.data.is_empty() {
|
||||
let nodes_json =
|
||||
serde_json::to_string(&data.data).unwrap_or_else(|_| "[]".to_string());
|
||||
let js = format!(
|
||||
r#"if (window.__loadAttackChain) {{ window.__loadAttackChain({nodes_json}); }}"#
|
||||
);
|
||||
document::eval(&js);
|
||||
// Small delay to ensure the DOM container exists
|
||||
spawn(async move {
|
||||
#[cfg(feature = "web")]
|
||||
gloo_timers::future::TimeoutFuture::new(100).await;
|
||||
#[cfg(not(feature = "web"))]
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
let js = format!(
|
||||
r#"if (window.__loadAttackChain) {{ window.__loadAttackChain({nodes_json}); }}"#
|
||||
);
|
||||
document::eval(&js);
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -109,7 +273,7 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
|
||||
let mut do_send_click = do_send.clone();
|
||||
|
||||
// Export handler
|
||||
// Export handlers
|
||||
let sid_for_export = session_id.clone();
|
||||
let do_export_md = move |_| {
|
||||
let sid = sid_for_export.clone();
|
||||
@@ -117,7 +281,6 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
spawn(async move {
|
||||
match export_pentest_report(sid.clone(), "markdown".to_string()).await {
|
||||
Ok(content) => {
|
||||
// Trigger download via JS
|
||||
let escaped = content
|
||||
.replace('\\', "\\\\")
|
||||
.replace('`', "\\`")
|
||||
@@ -256,7 +419,6 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
}
|
||||
}
|
||||
div { style: "display: flex; gap: 8px; align-items: center;",
|
||||
// Export buttons
|
||||
div { style: "display: flex; gap: 4px;",
|
||||
button {
|
||||
class: "btn btn-ghost",
|
||||
@@ -318,7 +480,7 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
let tool_name = msg.get("tool_name").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
||||
let tool_status = msg.get("tool_status").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
||||
|
||||
if msg_type == "tool_call" || msg_type == "tool_result" {
|
||||
if msg_type == "tool_call" || msg_type == "tool_result" || role == "tool_result" {
|
||||
let tool_icon_style = match tool_status.as_str() {
|
||||
"success" => "color: #16a34a;",
|
||||
"error" => "color: #dc2626;",
|
||||
@@ -358,6 +520,8 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Assistant message — render markdown
|
||||
let rendered_html = markdown_to_html(&content);
|
||||
rsx! {
|
||||
div {
|
||||
key: "{i}",
|
||||
@@ -367,8 +531,8 @@ pub fn PentestSessionPage(session_id: String) -> Element {
|
||||
Icon { icon: BsCpu, width: 14, height: 14 }
|
||||
}
|
||||
div {
|
||||
style: "max-width: 80%; padding: 10px 14px; background: var(--bg-tertiary); border-radius: 12px 12px 12px 2px; font-size: 0.9rem; line-height: 1.5; white-space: pre-wrap;",
|
||||
"{content}"
|
||||
style: "max-width: 80%; padding: 10px 14px; background: var(--bg-tertiary); border-radius: 12px 12px 12px 2px; font-size: 0.9rem; line-height: 1.5;",
|
||||
dangerous_inner_html: "{rendered_html}",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user