diff --git a/compliance-dashboard/src/pages/pentest_session.rs b/compliance-dashboard/src/pages/pentest_session.rs index f2bd7b5..f8f13fa 100644 --- a/compliance-dashboard/src/pages/pentest_session.rs +++ b/compliance-dashboard/src/pages/pentest_session.rs @@ -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(""); + in_code_block = false; + } else { + if in_list { + html.push_str(""); + in_list = false; + } + html.push_str("
");
+                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("");
+                in_list = false;
+            }
+            html.push_str("
"); + continue; + } + + // Lists + if trimmed.starts_with("- ") || trimmed.starts_with("* ") { + if !in_list { + html.push_str(""); + } + if in_code_block { + html.push_str("
"); + } + + 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_content}" + ); + 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!( + "{}{bold_content}{}", + &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}", } } }