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

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

View File

@@ -390,10 +390,13 @@ impl PentestOrchestrator {
)
.await;
// Build LLM-facing summary: strip large fields
// (screenshots, raw HTML) to save context window
let llm_data = summarize_tool_output(&result.data);
serde_json::json!({
"summary": result.summary,
"findings_count": findings_count,
"data": result.data,
"data": llm_data,
})
.to_string()
}
@@ -465,21 +468,61 @@ impl PentestOrchestrator {
.await;
}
// If cleanup_test_user is requested, append a cleanup instruction
// Clean up test user via identity provider API if requested
if session
.config
.as_ref()
.is_some_and(|c| c.auth.cleanup_test_user)
{
let cleanup_msg = PentestMessage::user(
session_id.clone(),
"Testing is complete. Now please clean up: navigate to the application and delete \
the test user account that was created during this session. Confirm once done."
.to_string(),
);
let _ = self.db.pentest_messages().insert_one(&cleanup_msg).await;
if let Some(ref test_user) = session.test_user {
let http = reqwest::Client::new();
// We need the AgentConfig — read from env since orchestrator doesn't hold it
let config = crate::config::load_config();
match config {
Ok(cfg) => {
match crate::pentest::cleanup::cleanup_test_user(test_user, &cfg, &http)
.await
{
Ok(true) => {
tracing::info!(
username = test_user.username.as_deref(),
"Test user cleaned up via provider API"
);
// Mark as cleaned up in DB
if let Some(sid) = session.id {
let _ = self
.db
.pentest_sessions()
.update_one(
doc! { "_id": sid },
doc! { "$set": { "test_user.cleaned_up": true } },
)
.await;
}
}
Ok(false) => {
tracing::info!(
"Test user cleanup skipped (no provider configured)"
);
}
Err(e) => {
tracing::warn!(error = %e, "Test user cleanup failed");
let _ = self.event_tx.send(PentestEvent::Error {
message: format!("Test user cleanup failed: {e}"),
});
}
}
}
Err(e) => {
tracing::warn!(error = %e, "Could not load config for cleanup");
}
}
}
}
// Clean up the persistent browser session for this pentest
compliance_dast::tools::browser::cleanup_browser_session(&session_id).await;
let _ = self.event_tx.send(PentestEvent::Complete {
summary: format!(
"Pentest complete. {} findings from {} tool invocations.",
@@ -490,3 +533,82 @@ impl PentestOrchestrator {
Ok(())
}
}
/// Strip large fields from tool output before sending to the LLM.
/// Screenshots, raw HTML, and other bulky data are replaced with short summaries.
/// The full data is still stored in the DB for the report.
fn summarize_tool_output(data: &serde_json::Value) -> serde_json::Value {
let Some(obj) = data.as_object() else {
return data.clone();
};
let mut summarized = serde_json::Map::new();
for (key, value) in obj {
match key.as_str() {
// Replace screenshot base64 with a placeholder
"screenshot_base64" => {
if let Some(s) = value.as_str() {
if !s.is_empty() {
summarized.insert(
key.clone(),
serde_json::Value::String(
"[screenshot captured and saved to report]".to_string(),
),
);
continue;
}
}
summarized.insert(key.clone(), value.clone());
}
// Truncate raw HTML content
"html" => {
if let Some(s) = value.as_str() {
if s.len() > 2000 {
summarized.insert(
key.clone(),
serde_json::Value::String(format!(
"{}... [truncated, {} chars total]",
&s[..2000],
s.len()
)),
);
continue;
}
}
summarized.insert(key.clone(), value.clone());
}
// Truncate page text
"text" if value.as_str().is_some_and(|s| s.len() > 1500) => {
let s = value.as_str().unwrap_or_default();
summarized.insert(
key.clone(),
serde_json::Value::String(format!("{}... [truncated]", &s[..1500])),
);
}
// Trim large arrays (e.g., "elements", "links", "inputs")
"elements" | "links" | "inputs" => {
if let Some(arr) = value.as_array() {
if arr.len() > 15 {
let mut trimmed: Vec<serde_json::Value> = arr[..15].to_vec();
trimmed.push(serde_json::json!(format!(
"... and {} more",
arr.len() - 15
)));
summarized.insert(key.clone(), serde_json::Value::Array(trimmed));
continue;
}
}
summarized.insert(key.clone(), value.clone());
}
// Recursively summarize nested objects (e.g., "page" in get_content)
_ if value.is_object() => {
summarized.insert(key.clone(), summarize_tool_output(value));
}
// Keep everything else as-is
_ => {
summarized.insert(key.clone(), value.clone());
}
}
}
serde_json::Value::Object(summarized)
}