feat: pentest onboarding — streaming, browser automation, reports, user cleanup (#16)

Complete pentest feature overhaul: SSE streaming, session-persistent browser tool (CDP), AES-256 credential encryption, auto-screenshots in reports, code-level remediation correlation, SAST triage chunking, context window optimization, test user cleanup (Keycloak/Auth0/Okta), wizard dropdowns, attack chain improvements, architecture docs with Mermaid diagrams.

Co-authored-by: Sharang Parnerkar <parnerkarsharang@gmail.com>
Reviewed-on: #16

compliance-agent/src/api/handlers/pentest_handlers/session.rs

@@ -17,10 +17,12 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
 
 #[derive(Deserialize)]
 pub struct CreateSessionRequest {
-    pub target_id: String,
+    pub target_id: Option<String>,
     #[serde(default = "default_strategy")]
     pub strategy: String,
     pub message: Option<String>,
+    /// Wizard configuration — if present, takes precedence over legacy fields
+    pub config: Option<PentestConfig>,
 }
 
 fn default_strategy() -> String {
@@ -32,83 +34,365 @@ pub struct SendMessageRequest {
     pub message: String,
 }
 
+#[derive(Deserialize)]
+pub struct LookupRepoQuery {
+    pub url: String,
+}
+
 /// POST /api/v1/pentest/sessions — Create a new pentest session and start the orchestrator
 #[tracing::instrument(skip_all)]
 pub async fn create_session(
     Extension(agent): AgentExt,
     Json(req): Json<CreateSessionRequest>,
 ) -> Result<Json<ApiResponse<PentestSession>>, (StatusCode, String)> {
-    let oid = mongodb::bson::oid::ObjectId::parse_str(&req.target_id).map_err(|_| {
-        (
-            StatusCode::BAD_REQUEST,
-            "Invalid target_id format".to_string(),
-        )
-    })?;
-
-    // Look up the target
-    let target = agent
-        .db
-        .dast_targets()
-        .find_one(doc! { "_id": oid })
-        .await
-        .map_err(|e| {
+    // Try to acquire a concurrency permit
+    let permit = agent
+        .session_semaphore
+        .clone()
+        .try_acquire_owned()
+        .map_err(|_| {
             (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                format!("Database error: {e}"),
+                StatusCode::TOO_MANY_REQUESTS,
+                "Maximum concurrent pentest sessions reached. Try again later.".to_string(),
             )
-        })?
-        .ok_or_else(|| (StatusCode::NOT_FOUND, "Target not found".to_string()))?;
+        })?;
 
-    // Parse strategy
-    let strategy = match req.strategy.as_str() {
+    if let Some(ref config) = req.config {
+        // ── Wizard path ──────────────────────────────────────────────
+        if !config.disclaimer_accepted {
+            return Err((
+                StatusCode::BAD_REQUEST,
+                "Disclaimer must be accepted".to_string(),
+            ));
+        }
+
+        // Look up or auto-create DastTarget by app_url
+        let target = match agent
+            .db
+            .dast_targets()
+            .find_one(doc! { "base_url": &config.app_url })
+            .await
+            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?
+        {
+            Some(t) => t,
+            None => {
+                use compliance_core::models::dast::{DastTarget, DastTargetType};
+                let mut t = DastTarget::new(
+                    config.app_url.clone(),
+                    config.app_url.clone(),
+                    DastTargetType::WebApp,
+                );
+                if let Some(rl) = config.rate_limit {
+                    t.rate_limit = rl;
+                }
+                t.allow_destructive = config.allow_destructive;
+                t.excluded_paths = config.scope_exclusions.clone();
+                let res = agent.db.dast_targets().insert_one(&t).await.map_err(|e| {
+                    (
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        format!("Failed to create target: {e}"),
+                    )
+                })?;
+                t.id = res.inserted_id.as_object_id();
+                t
+            }
+        };
+
+        let target_id = target.id.map(|oid| oid.to_hex()).unwrap_or_default();
+
+        // Parse strategy from config or request
+        let strat_str = config.strategy.as_deref().unwrap_or(req.strategy.as_str());
+        let strategy = parse_strategy(strat_str);
+
+        let mut session = PentestSession::new(target_id, strategy);
+        session.config = Some(config.clone());
+        session.repo_id = target.repo_id.clone();
+
+        // Resolve repo_id from git_repo_url if provided
+        if let Some(ref git_url) = config.git_repo_url {
+            if let Ok(Some(repo)) = agent
+                .db
+                .repositories()
+                .find_one(doc! { "git_url": git_url })
+                .await
+            {
+                session.repo_id = repo.id.map(|oid| oid.to_hex());
+            }
+        }
+
+        let insert_result = agent
+            .db
+            .pentest_sessions()
+            .insert_one(&session)
+            .await
+            .map_err(|e| {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Failed to create session: {e}"),
+                )
+            })?;
+        session.id = insert_result.inserted_id.as_object_id();
+
+        let session_id_str = session.id.map(|oid| oid.to_hex()).unwrap_or_default();
+
+        // Register broadcast stream and pause control
+        let event_tx = agent.register_session_stream(&session_id_str);
+        let pause_rx = agent.register_pause_control(&session_id_str);
+
+        // Merge server-default IMAP/email settings where wizard left blanks
+        if let Some(ref mut cfg) = session.config {
+            if cfg.auth.mode == AuthMode::AutoRegister {
+                if cfg.auth.verification_email.is_none() {
+                    cfg.auth.verification_email = agent.config.pentest_verification_email.clone();
+                }
+                if cfg.auth.imap_host.is_none() {
+                    cfg.auth.imap_host = agent.config.pentest_imap_host.clone();
+                }
+                if cfg.auth.imap_port.is_none() {
+                    cfg.auth.imap_port = agent.config.pentest_imap_port;
+                }
+                if cfg.auth.imap_username.is_none() {
+                    cfg.auth.imap_username = agent.config.pentest_imap_username.clone();
+                }
+                if cfg.auth.imap_password.is_none() {
+                    cfg.auth.imap_password = agent.config.pentest_imap_password.as_ref().map(|s| {
+                        use secrecy::ExposeSecret;
+                        s.expose_secret().to_string()
+                    });
+                }
+            }
+        }
+
+        // Pre-populate test user record for auto-register sessions
+        if let Some(ref cfg) = session.config {
+            if cfg.auth.mode == AuthMode::AutoRegister {
+                let verification_email = cfg.auth.verification_email.clone();
+                // Build plus-addressed email for this session
+                let test_email = verification_email.as_deref().map(|email| {
+                    let parts: Vec<&str> = email.splitn(2, '@').collect();
+                    if parts.len() == 2 {
+                        format!("{}+{}@{}", parts[0], session_id_str, parts[1])
+                    } else {
+                        email.to_string()
+                    }
+                });
+
+                // Detect identity provider from keycloak config
+                let provider = if agent.config.keycloak_url.is_some() {
+                    Some(compliance_core::models::pentest::IdentityProvider::Keycloak)
+                } else {
+                    None
+                };
+
+                session.test_user = Some(compliance_core::models::pentest::TestUserRecord {
+                    username: None, // LLM will choose; updated after registration
+                    email: test_email,
+                    provider_user_id: None,
+                    provider,
+                    cleaned_up: false,
+                });
+            }
+        }
+
+        // Encrypt credentials before they linger in memory
+        let mut session_for_task = session.clone();
+        if let Some(ref mut cfg) = session_for_task.config {
+            cfg.auth.username = cfg
+                .auth
+                .username
+                .as_ref()
+                .map(|u| crate::pentest::crypto::encrypt(u));
+            cfg.auth.password = cfg
+                .auth
+                .password
+                .as_ref()
+                .map(|p| crate::pentest::crypto::encrypt(p));
+        }
+
+        // Persist encrypted credentials to DB
+        if session_for_task.config.is_some() {
+            if let Some(sid) = session.id {
+                let _ = agent
+                    .db
+                    .pentest_sessions()
+                    .update_one(
+                        doc! { "_id": sid },
+                        doc! { "$set": {
+                            "config.auth.username": session_for_task.config.as_ref()
+                                .and_then(|c| c.auth.username.as_deref())
+                                .map(|s| mongodb::bson::Bson::String(s.to_string()))
+                                .unwrap_or(mongodb::bson::Bson::Null),
+                            "config.auth.password": session_for_task.config.as_ref()
+                                .and_then(|c| c.auth.password.as_deref())
+                                .map(|s| mongodb::bson::Bson::String(s.to_string()))
+                                .unwrap_or(mongodb::bson::Bson::Null),
+                        }},
+                    )
+                    .await;
+            }
+        }
+
+        let initial_message = config
+            .initial_instructions
+            .clone()
+            .or(req.message.clone())
+            .unwrap_or_else(|| {
+                format!(
+                    "Begin a {} penetration test against {} ({}). \
+                     Identify vulnerabilities and provide evidence for each finding.",
+                    session.strategy, target.name, target.base_url,
+                )
+            });
+
+        let llm = agent.llm.clone();
+        let db = agent.db.clone();
+        let session_clone = session.clone();
+        let target_clone = target.clone();
+        let agent_ref = agent.clone();
+        tokio::spawn(async move {
+            let orchestrator = PentestOrchestrator::new(llm, db, event_tx, Some(pause_rx));
+            orchestrator
+                .run_session_guarded(&session_clone, &target_clone, &initial_message)
+                .await;
+            // Clean up session resources
+            agent_ref.cleanup_session(&session_id_str);
+            // Release concurrency permit
+            drop(permit);
+        });
+
+        // Redact credentials in response
+        let mut response_session = session;
+        if let Some(ref mut cfg) = response_session.config {
+            if cfg.auth.username.is_some() {
+                cfg.auth.username = Some("********".to_string());
+            }
+            if cfg.auth.password.is_some() {
+                cfg.auth.password = Some("********".to_string());
+            }
+        }
+
+        Ok(Json(ApiResponse {
+            data: response_session,
+            total: None,
+            page: None,
+        }))
+    } else {
+        // ── Legacy path ──────────────────────────────────────────────
+        let target_id = req.target_id.clone().ok_or_else(|| {
+            (
+                StatusCode::BAD_REQUEST,
+                "target_id is required for legacy creation".to_string(),
+            )
+        })?;
+
+        let oid = mongodb::bson::oid::ObjectId::parse_str(&target_id).map_err(|_| {
+            (
+                StatusCode::BAD_REQUEST,
+                "Invalid target_id format".to_string(),
+            )
+        })?;
+
+        let target = agent
+            .db
+            .dast_targets()
+            .find_one(doc! { "_id": oid })
+            .await
+            .map_err(|e| {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Database error: {e}"),
+                )
+            })?
+            .ok_or_else(|| (StatusCode::NOT_FOUND, "Target not found".to_string()))?;
+
+        let strategy = parse_strategy(&req.strategy);
+
+        let mut session = PentestSession::new(target_id, strategy);
+        session.repo_id = target.repo_id.clone();
+
+        let insert_result = agent
+            .db
+            .pentest_sessions()
+            .insert_one(&session)
+            .await
+            .map_err(|e| {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Failed to create session: {e}"),
+                )
+            })?;
+        session.id = insert_result.inserted_id.as_object_id();
+
+        let session_id_str = session.id.map(|oid| oid.to_hex()).unwrap_or_default();
+
+        // Register broadcast stream and pause control
+        let event_tx = agent.register_session_stream(&session_id_str);
+        let pause_rx = agent.register_pause_control(&session_id_str);
+
+        let initial_message = req.message.unwrap_or_else(|| {
+            format!(
+                "Begin a {} penetration test against {} ({}). \
+                 Identify vulnerabilities and provide evidence for each finding.",
+                session.strategy, target.name, target.base_url,
+            )
+        });
+
+        let llm = agent.llm.clone();
+        let db = agent.db.clone();
+        let session_clone = session.clone();
+        let target_clone = target.clone();
+        let agent_ref = agent.clone();
+        tokio::spawn(async move {
+            let orchestrator = PentestOrchestrator::new(llm, db, event_tx, Some(pause_rx));
+            orchestrator
+                .run_session_guarded(&session_clone, &target_clone, &initial_message)
+                .await;
+            agent_ref.cleanup_session(&session_id_str);
+            drop(permit);
+        });
+
+        Ok(Json(ApiResponse {
+            data: session,
+            total: None,
+            page: None,
+        }))
+    }
+}
+
+fn parse_strategy(s: &str) -> PentestStrategy {
+    match s {
         "quick" => PentestStrategy::Quick,
         "targeted" => PentestStrategy::Targeted,
         "aggressive" => PentestStrategy::Aggressive,
         "stealth" => PentestStrategy::Stealth,
         _ => PentestStrategy::Comprehensive,
+    }
+}
+
+/// GET /api/v1/pentest/lookup-repo — Look up a tracked repository by git URL
+#[tracing::instrument(skip_all)]
+pub async fn lookup_repo(
+    Extension(agent): AgentExt,
+    Query(params): Query<LookupRepoQuery>,
+) -> Result<Json<ApiResponse<serde_json::Value>>, StatusCode> {
+    let repo = agent
+        .db
+        .repositories()
+        .find_one(doc! { "git_url": &params.url })
+        .await
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    let data = match repo {
+        Some(r) => serde_json::json!({
+            "name": r.name,
+            "default_branch": r.default_branch,
+            "last_scanned_commit": r.last_scanned_commit,
+        }),
+        None => serde_json::Value::Null,
     };
 
-    // Create session
-    let mut session = PentestSession::new(req.target_id.clone(), strategy);
-    session.repo_id = target.repo_id.clone();
-
-    let insert_result = agent
-        .db
-        .pentest_sessions()
-        .insert_one(&session)
-        .await
-        .map_err(|e| {
-            (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                format!("Failed to create session: {e}"),
-            )
-        })?;
-
-    // Set the generated ID back on the session so the orchestrator has it
-    session.id = insert_result.inserted_id.as_object_id();
-
-    let initial_message = req.message.unwrap_or_else(|| {
-        format!(
-            "Begin a {} penetration test against {} ({}). \
-             Identify vulnerabilities and provide evidence for each finding.",
-            session.strategy, target.name, target.base_url,
-        )
-    });
-
-    // Spawn the orchestrator on a background task
-    let llm = agent.llm.clone();
-    let db = agent.db.clone();
-    let session_clone = session.clone();
-    let target_clone = target.clone();
-    tokio::spawn(async move {
-        let orchestrator = PentestOrchestrator::new(llm, db);
-        orchestrator
-            .run_session_guarded(&session_clone, &target_clone, &initial_message)
-            .await;
-    });
-
     Ok(Json(ApiResponse {
-        data: session,
+        data,
         total: None,
         page: None,
     }))
@@ -158,7 +442,7 @@ pub async fn get_session(
 ) -> Result<Json<ApiResponse<PentestSession>>, StatusCode> {
     let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
 
-    let session = agent
+    let mut session = agent
         .db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
@@ -166,6 +450,16 @@ pub async fn get_session(
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
         .ok_or(StatusCode::NOT_FOUND)?;
 
+    // Redact credentials in response
+    if let Some(ref mut cfg) = session.config {
+        if cfg.auth.username.is_some() {
+            cfg.auth.username = Some("********".to_string());
+        }
+        if cfg.auth.password.is_some() {
+            cfg.auth.password = Some("********".to_string());
+        }
+    }
+
     Ok(Json(ApiResponse {
         data: session,
         total: None,
@@ -241,8 +535,20 @@ pub async fn send_message(
     let llm = agent.llm.clone();
     let db = agent.db.clone();
     let message = req.message.clone();
+
+    // Use existing broadcast sender if available, otherwise create a new one
+    let event_tx = agent
+        .subscribe_session(&session_id)
+        .and_then(|_| {
+            agent
+                .session_streams
+                .get(&session_id)
+                .map(|entry| entry.value().clone())
+        })
+        .unwrap_or_else(|| agent.register_session_stream(&session_id));
+
     tokio::spawn(async move {
-        let orchestrator = PentestOrchestrator::new(llm, db);
+        let orchestrator = PentestOrchestrator::new(llm, db, event_tx, None);
         orchestrator
             .run_session_guarded(&session, &target, &message)
             .await;
@@ -277,10 +583,10 @@ pub async fn stop_session(
         })?
         .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?;
 
-    if session.status != PentestStatus::Running {
+    if session.status != PentestStatus::Running && session.status != PentestStatus::Paused {
         return Err((
             StatusCode::BAD_REQUEST,
-            format!("Session is {}, not running", session.status),
+            format!("Session is {}, not running or paused", session.status),
         ));
     }
 
@@ -303,6 +609,9 @@ pub async fn stop_session(
             )
         })?;
 
+    // Clean up session resources
+    agent.cleanup_session(&id);
+
     let updated = agent
         .db
         .pentest_sessions()
@@ -328,6 +637,92 @@ pub async fn stop_session(
     }))
 }
 
+/// POST /api/v1/pentest/sessions/:id/pause — Pause a running pentest session
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn pause_session(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+) -> Result<Json<ApiResponse<serde_json::Value>>, (StatusCode, String)> {
+    let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
+        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
+
+    let session = agent
+        .db
+        .pentest_sessions()
+        .find_one(doc! { "_id": oid })
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Database error: {e}"),
+            )
+        })?
+        .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?;
+
+    if session.status != PentestStatus::Running {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            format!("Session is {}, not running", session.status),
+        ));
+    }
+
+    if !agent.pause_session(&id) {
+        return Err((
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "Failed to send pause signal".to_string(),
+        ));
+    }
+
+    Ok(Json(ApiResponse {
+        data: serde_json::json!({ "status": "paused" }),
+        total: None,
+        page: None,
+    }))
+}
+
+/// POST /api/v1/pentest/sessions/:id/resume — Resume a paused pentest session
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn resume_session(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+) -> Result<Json<ApiResponse<serde_json::Value>>, (StatusCode, String)> {
+    let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
+        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
+
+    let session = agent
+        .db
+        .pentest_sessions()
+        .find_one(doc! { "_id": oid })
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Database error: {e}"),
+            )
+        })?
+        .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?;
+
+    if session.status != PentestStatus::Paused {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            format!("Session is {}, not paused", session.status),
+        ));
+    }
+
+    if !agent.resume_session(&id) {
+        return Err((
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "Failed to send resume signal".to_string(),
+        ));
+    }
+
+    Ok(Json(ApiResponse {
+        data: serde_json::json!({ "status": "running" }),
+        total: None,
+        page: None,
+    }))
+}
+
 /// GET /api/v1/pentest/sessions/:id/attack-chain — Get attack chain nodes for a session
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn get_attack_chain(