sharang/compliance-scanner-agent
1
0
Fork 0
Code Issues 51 Pull Requests 1 Actions Packages Projects Releases 1 Wiki Activity

feat: browser session persistence, auto-screenshots, context optimization, user cleanup
Some checks failed
CI / Check (pull_request) Failing after 5m55s
Details
CI / Detect Changes (pull_request) Has been skipped
Details
CI / Deploy Agent (pull_request) Has been skipped
Details
CI / Deploy Dashboard (pull_request) Has been skipped
Details
CI / Deploy Docs (pull_request) Has been skipped
Details
CI / Deploy MCP (pull_request) Has been skipped
Details

Browse Source 
Browser tool:
- Session-persistent Chrome tab (same tab reused across all calls in a pentest)
- Auto-screenshot on every navigate and click (stored in attack chain for report)
- Fill uses CDP Input.insertText (fixes WebSocket corruption on special chars)
- Switched from browserless/chromium to chromedp/headless-shell (stable WS)

Context window optimization:
- Strip screenshot_base64 from LLM conversation (kept in DB for report)
- Truncate HTML to 2KB, page text to 1.5KB in LLM messages
- Cap element/link arrays at 15 items
- SAST triage: batch 30 findings per LLM call instead of all at once

Report improvements:
- Auto-embed screenshots in attack chain timeline (navigate + click nodes)
- Cover page shows best app screenshot
- Attack chain phases capped at 8 (no more 20x "Final")

User cleanup:
- TestUserRecord model tracks created test users per session
- cleanup.rs: Keycloak (Admin REST API), Auth0 (Management API), Okta (Users API)
- Auto-cleanup on session completion when cleanup_test_user is enabled
- Env vars: KEYCLOAK_ADMIN_USERNAME, KEYCLOAK_ADMIN_PASSWORD

System prompt:
- Explicit browser usage instructions (navigate → get_content → click → fill)
- SPA auth bypass guidance (check page content, not HTTP status)
- Screenshot instructions for evidence collection

Other:
- Pin mongo:7 in docker-compose (mongo:latest/8 segfaults on kernel 6.19)
- Add deploy/docker-compose.mailserver.yml for Postfix + Dovecot

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sharang Parnerkar
2026-03-17 19:53:55 +01:00
parent a737c36bc9
commit 37690ce734
18 changed files with 1122 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files

350
compliance-dast/src/tools/browser.rs
View File

@@ -1,4 +1,6 @@
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use base64::Engine;
@@ -6,17 +8,26 @@ use compliance_core::error::CoreError;
use compliance_core::traits::pentest_tool::{PentestTool, PentestToolContext, PentestToolResult};
use futures_util::{SinkExt, StreamExt};
use serde_json::json;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;
use tracing::info;
type WsStream =
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
/// Global pool of persistent browser sessions keyed by pentest session ID.
/// Each pentest session gets one Chrome tab that stays alive across tool calls.
static BROWSER_SESSIONS: std::sync::LazyLock<Arc<Mutex<HashMap<String, BrowserSession>>>> =
std::sync::LazyLock::new(|| Arc::new(Mutex::new(HashMap::new())));
/// A browser automation tool that exposes headless Chrome actions to the LLM
/// via the Chrome DevTools Protocol. Reuses the same `CHROME_WS_URL` used for
/// PDF generation.
/// via the Chrome DevTools Protocol.
///
/// Supported actions: navigate, screenshot, click, fill, get_content, evaluate.
/// **Session-persistent**: the same Chrome tab is reused across all invocations
/// within a pentest session, so cookies, auth state, and page context are
/// preserved between navigate → click → fill → screenshot calls.
///
/// Supported actions: navigate, screenshot, click, fill, get_content, evaluate, close.
pub struct BrowserTool;
impl Default for BrowserTool {
@@ -31,11 +42,13 @@ impl PentestTool for BrowserTool {
}
fn description(&self) -> &str {
"Headless browser automation via Chrome DevTools Protocol. \
Supports navigating to URLs, taking screenshots, clicking elements, \
"Headless browser automation via Chrome DevTools Protocol. The browser tab persists \
across calls within the same pentest session — cookies, login state, and page context \
are preserved. Supports navigating to URLs, taking screenshots, clicking elements, \
filling form fields, reading page content, and evaluating JavaScript. \
Use CSS selectors to target elements. Useful for discovering registration pages, \
filling out forms, extracting verification links, and visual inspection."
Use CSS selectors to target elements. After navigating, use get_content to read the \
page HTML and find elements to click or fill. Use this to discover registration pages, \
fill out signup forms, complete email verification, and test authenticated flows."
}
fn input_schema(&self) -> serde_json::Value {
@@ -44,8 +57,8 @@ impl PentestTool for BrowserTool {
"properties": {
"action": {
"type": "string",
"enum": ["navigate", "screenshot", "click", "fill", "get_content", "evaluate"],
"description": "Action to perform"
"enum": ["navigate", "screenshot", "click", "fill", "get_content", "evaluate", "close"],
"description": "Action to perform. The browser tab persists between calls — use navigate first, then get_content to see the page, then click/fill to interact."
},
"url": {
"type": "string",
@@ -53,7 +66,7 @@ impl PentestTool for BrowserTool {
},
"selector": {
"type": "string",
"description": "CSS selector for click/fill actions"
"description": "CSS selector for click/fill actions (e.g. '#username', 'a[href*=register]', 'button[type=submit]')"
},
"value": {
"type": "string",
@@ -61,7 +74,7 @@ impl PentestTool for BrowserTool {
},
"wait_ms": {
"type": "integer",
"description": "Milliseconds to wait after action (default: 500)"
"description": "Milliseconds to wait after action (default: 1000)"
}
},
"required": ["action"]
@@ -71,7 +84,7 @@ impl PentestTool for BrowserTool {
fn execute<'a>(
&'a self,
input: serde_json::Value,
_context: &'a PentestToolContext,
context: &'a PentestToolContext,
) -> Pin<Box<dyn std::future::Future<Output = Result<PentestToolResult, CoreError>> + Send + 'a>>
{
Box::pin(async move {
@@ -79,11 +92,42 @@ impl PentestTool for BrowserTool {
let url = input.get("url").and_then(|v| v.as_str()).unwrap_or("");
let selector = input.get("selector").and_then(|v| v.as_str()).unwrap_or("");
let value = input.get("value").and_then(|v| v.as_str()).unwrap_or("");
let wait_ms = input.get("wait_ms").and_then(|v| v.as_u64()).unwrap_or(500);
let wait_ms = input
.get("wait_ms")
.and_then(|v| v.as_u64())
.unwrap_or(1000);
let session_key = context.session_id.clone();
let mut session = BrowserSession::connect()
.await
.map_err(|e| CoreError::Other(format!("Browser connect failed: {e}")))?;
// Handle close action — tear down the persistent session
if action == "close" {
let mut pool = BROWSER_SESSIONS.lock().await;
if let Some(mut sess) = pool.remove(&session_key) {
let _ = sess.close().await;
}
return Ok(PentestToolResult {
summary: "Browser session closed".to_string(),
findings: Vec::new(),
data: json!({ "closed": true }),
});
}
// Get or create persistent session for this pentest
let mut pool = BROWSER_SESSIONS.lock().await;
if !pool.contains_key(&session_key) {
match BrowserSession::connect().await {
Ok(sess) => {
pool.insert(session_key.clone(), sess);
}
Err(e) => {
return Err(CoreError::Other(format!("Browser connect failed: {e}")));
}
}
}
let session = pool.get_mut(&session_key);
let Some(session) = session else {
return Err(CoreError::Other("Browser session not found".to_string()));
};
let result = match action {
"navigate" => session.navigate(url, wait_ms).await,
@@ -95,8 +139,15 @@ impl PentestTool for BrowserTool {
_ => Err(format!("Unknown browser action: {action}")),
};
// Always try to clean up
let _ = session.close().await;
// If the session errored, remove it so the next call creates a fresh one
if result.is_err() {
if let Some(mut dead) = pool.remove(&session_key) {
let _ = dead.close().await;
}
}
// Release the lock before building the response
drop(pool);
match result {
Ok(data) => {
@@ -214,7 +265,7 @@ impl BrowserSession {
}
async fn navigate(&mut self, url: &str, wait_ms: u64) -> Result<serde_json::Value, String> {
let resp = cdp_send_session(
cdp_send_session(
&mut self.ws,
self.next_id,
&self.session_id,
@@ -226,19 +277,44 @@ impl BrowserSession {
tokio::time::sleep(Duration::from_millis(wait_ms)).await;
// Get page title
let title_resp = self.evaluate_raw("document.title").await?;
let page_url_resp = self.evaluate_raw("window.location.href").await?;
// Get page title and current URL (may have redirected)
let title = self
.evaluate_raw("document.title")
.await
.unwrap_or_default();
let page_url = self
.evaluate_raw("window.location.href")
.await
.unwrap_or_default();
// Auto-get a summary of interactive elements on the page
let links_js = r#"(function(){
var items = [];
document.querySelectorAll('a[href]').forEach(function(a, i) {
if (i < 20) items.push({tag:'a', text:a.textContent.trim().substring(0,60), href:a.href});
});
document.querySelectorAll('input,select,textarea,button[type=submit]').forEach(function(el, i) {
if (i < 20) items.push({tag:el.tagName.toLowerCase(), type:el.type||'', name:el.name||'', id:el.id||'', placeholder:el.placeholder||''});
});
return JSON.stringify(items);
})()"#;
let elements_json = self.evaluate_raw(links_js).await.unwrap_or_default();
let elements: serde_json::Value = serde_json::from_str(&elements_json).unwrap_or(json!([]));
// Auto-capture screenshot after every navigation
let screenshot_b64 = self.capture_screenshot_b64().await.unwrap_or_default();
Ok(json!({
"navigated": true,
"url": page_url_resp,
"title": title_resp,
"frame_id": resp.get("result").and_then(|r| r.get("frameId")),
"url": page_url,
"title": title,
"elements": elements,
"screenshot_base64": screenshot_b64,
}))
}
async fn screenshot(&mut self) -> Result<serde_json::Value, String> {
/// Capture a screenshot and return the base64 string (empty on failure).
async fn capture_screenshot_b64(&mut self) -> Result<String, String> {
let resp = cdp_send_session(
&mut self.ws,
self.next_id,
@@ -249,14 +325,19 @@ impl BrowserSession {
.await?;
self.next_id += 1;
let b64 = resp
Ok(resp
.get("result")
.and_then(|r| r.get("data"))
.and_then(|d| d.as_str())
.unwrap_or("");
.unwrap_or("")
.to_string())
}
async fn screenshot(&mut self) -> Result<serde_json::Value, String> {
let b64 = self.capture_screenshot_b64().await?;
let size_kb = base64::engine::general_purpose::STANDARD
.decode(b64)
.decode(&b64)
.map(|b| b.len() / 1024)
.unwrap_or(0);
@@ -267,7 +348,6 @@ impl BrowserSession {
}
async fn click(&mut self, selector: &str, wait_ms: u64) -> Result<serde_json::Value, String> {
// Use JS to find element and get its bounding box, then click
let js = format!(
r#"(function() {{
var el = document.querySelector({sel});
@@ -289,9 +369,29 @@ impl BrowserSession {
let result = self.evaluate_raw(&js).await?;
tokio::time::sleep(Duration::from_millis(wait_ms)).await;
serde_json::from_str::<serde_json::Value>(&result)
.unwrap_or_else(|_| json!({ "result": result }));
Ok(serde_json::from_str(&result).unwrap_or(json!({ "result": result })))
// After click, get current URL (may have navigated)
let current_url = self
.evaluate_raw("window.location.href")
.await
.unwrap_or_default();
let title = self
.evaluate_raw("document.title")
.await
.unwrap_or_default();
// Auto-capture screenshot after click
let screenshot_b64 = self.capture_screenshot_b64().await.unwrap_or_default();
let mut click_result: serde_json::Value =
serde_json::from_str(&result).unwrap_or(json!({ "result": result }));
if let Some(obj) = click_result.as_object_mut() {
obj.insert("current_url".to_string(), json!(current_url));
obj.insert("page_title".to_string(), json!(title));
if !screenshot_b64.is_empty() {
obj.insert("screenshot_base64".to_string(), json!(screenshot_b64));
}
}
Ok(click_result)
}
async fn fill(
@@ -300,62 +400,83 @@ impl BrowserSession {
value: &str,
wait_ms: u64,
) -> Result<serde_json::Value, String> {
let js = format!(
r#"(function() {{
var el = document.querySelector({sel});
if (!el) return JSON.stringify({{error: "Element not found: {raw}"}});
el.focus();
el.value = {val};
el.dispatchEvent(new Event('input', {{bubbles: true}}));
el.dispatchEvent(new Event('change', {{bubbles: true}}));
return JSON.stringify({{filled: true, tag: el.tagName, selector: {sel}}});
}})()"#,
// Step 1: Focus the element via JS
let focus_js = format!(
"(function(){{var e=document.querySelector({sel});\
if(!e)return 'notfound';e.focus();e.select();return 'ok'}})()",
sel = serde_json::to_string(selector).unwrap_or_default(),
raw = selector.replace('"', r#"\""#),
val = serde_json::to_string(value).unwrap_or_default(),
);
let found = self.evaluate_raw(&focus_js).await?;
if found == "notfound" {
return Ok(json!({ "error": format!("Element not found: {selector}") }));
}
// Step 2: Clear existing content with Select All + Delete
cdp_send_session(
&mut self.ws,
self.next_id,
&self.session_id,
"Input.dispatchKeyEvent",
json!({"type": "keyDown", "key": "a", "code": "KeyA", "modifiers": 2}),
)
.await?;
self.next_id += 1;
cdp_send_session(
&mut self.ws,
self.next_id,
&self.session_id,
"Input.dispatchKeyEvent",
json!({"type": "keyUp", "key": "a", "code": "KeyA", "modifiers": 2}),
)
.await?;
self.next_id += 1;
cdp_send_session(
&mut self.ws,
self.next_id,
&self.session_id,
"Input.dispatchKeyEvent",
json!({"type": "keyDown", "key": "Backspace", "code": "Backspace"}),
)
.await?;
self.next_id += 1;
cdp_send_session(
&mut self.ws,
self.next_id,
&self.session_id,
"Input.dispatchKeyEvent",
json!({"type": "keyUp", "key": "Backspace", "code": "Backspace"}),
)
.await?;
self.next_id += 1;
// Step 3: Insert the text using Input.insertText (single CDP command, no JS eval)
cdp_send_session(
&mut self.ws,
self.next_id,
&self.session_id,
"Input.insertText",
json!({"text": value}),
)
.await?;
self.next_id += 1;
// Step 4: Verify the value was set
let verify_js = format!(
"(function(){{var e=document.querySelector({sel});return e?e.value:''}})()",
sel = serde_json::to_string(selector).unwrap_or_default(),
);
let final_value = self.evaluate_raw(&verify_js).await.unwrap_or_default();
let result = self.evaluate_raw(&js).await?;
tokio::time::sleep(Duration::from_millis(wait_ms)).await;
Ok(serde_json::from_str(&result).unwrap_or(json!({ "result": result })))
Ok(json!({
"filled": true,
"selector": selector,
"value": final_value,
}))
}
async fn get_content(&mut self) -> Result<serde_json::Value, String> {
let resp = cdp_send_session(
&mut self.ws,
self.next_id,
&self.session_id,
"DOM.getDocument",
json!({ "depth": 0 }),
)
.await?;
self.next_id += 1;
let root_id = resp
.get("result")
.and_then(|r| r.get("root"))
.and_then(|n| n.get("nodeId"))
.and_then(|n| n.as_i64())
.unwrap_or(1);
let html_resp = cdp_send_session(
&mut self.ws,
self.next_id,
&self.session_id,
"DOM.getOuterHTML",
json!({ "nodeId": root_id }),
)
.await?;
self.next_id += 1;
let html = html_resp
.get("result")
.and_then(|r| r.get("outerHTML"))
.and_then(|h| h.as_str())
.unwrap_or("");
// Also get page title and URL for context
let title = self
.evaluate_raw("document.title")
.await
@@ -365,22 +486,55 @@ impl BrowserSession {
.await
.unwrap_or_default();
// Truncate HTML to avoid massive payloads to the LLM
let truncated = if html.len() > 50_000 {
format!(
"{}... [truncated, {} total chars]",
&html[..50_000],
html.len()
)
} else {
html.to_string()
};
// Get a structured summary instead of raw HTML (more useful for LLM)
let summary_js = r#"(function(){
var result = {forms:[], links:[], inputs:[], buttons:[], headings:[], text:''};
// Forms
document.querySelectorAll('form').forEach(function(f,i){
if(i<10) result.forms.push({action:f.action, method:f.method, id:f.id});
});
// Links
document.querySelectorAll('a[href]').forEach(function(a,i){
if(i<30) result.links.push({text:a.textContent.trim().substring(0,80), href:a.href});
});
// Inputs
document.querySelectorAll('input,select,textarea').forEach(function(el,i){
if(i<30) result.inputs.push({
tag:el.tagName.toLowerCase(),
type:el.type||'',
name:el.name||'',
id:el.id||'',
placeholder:el.placeholder||'',
value:el.type==='password'?'***':el.value.substring(0,50)
});
});
// Buttons
document.querySelectorAll('button,[type=submit],[role=button]').forEach(function(b,i){
if(i<20) result.buttons.push({text:b.textContent.trim().substring(0,60), type:b.type||'', id:b.id||''});
});
// Headings
document.querySelectorAll('h1,h2,h3').forEach(function(h,i){
if(i<10) result.headings.push(h.textContent.trim().substring(0,100));
});
// Page text (truncated)
result.text = document.body ? document.body.innerText.substring(0, 3000) : '';
return JSON.stringify(result);
})()"#;
let summary = self.evaluate_raw(summary_js).await.unwrap_or_default();
let page_data: serde_json::Value = serde_json::from_str(&summary).unwrap_or(json!({}));
Ok(json!({
"url": url,
"title": title,
"html": truncated,
"html_length": html.len(),
"page": page_data,
}))
}
@@ -431,7 +585,15 @@ impl BrowserSession {
}
}
// ── CDP helpers (same pattern as compliance-agent/src/pentest/report/pdf.rs) ──
/// Clean up the browser session for a pentest session (call when session ends).
pub async fn cleanup_browser_session(session_id: &str) {
let mut pool = BROWSER_SESSIONS.lock().await;
if let Some(mut sess) = pool.remove(session_id) {
let _ = sess.close().await;
}
}
// ── CDP helpers ──
async fn cdp_send(
ws: &mut WsStream,