compliance-scanner-agent/compliance-agent/src/pentest/cleanup.rs

use compliance_core::models::pentest::{IdentityProvider, TestUserRecord};
use compliance_core::AgentConfig;
use secrecy::ExposeSecret;
use tracing::{info, warn};

/// Attempt to delete a test user created during a pentest session.
///
/// Routes to the appropriate identity provider based on `TestUserRecord.provider`.
/// Falls back to browser-based cleanup if no API credentials are available.
///
/// Returns `Ok(true)` if the user was deleted, `Ok(false)` if skipped, `Err` on failure.
pub async fn cleanup_test_user(
    user: &TestUserRecord,
    config: &AgentConfig,
    http: &reqwest::Client,
) -> Result<bool, String> {
    if user.cleaned_up {
        return Ok(false);
    }

    let provider = user.provider.as_ref();

    match provider {
        Some(IdentityProvider::Keycloak) => cleanup_keycloak(user, config, http).await,
        Some(IdentityProvider::Auth0) => cleanup_auth0(user, config, http).await,
        Some(IdentityProvider::Okta) => cleanup_okta(user, config, http).await,
        Some(IdentityProvider::Firebase) => {
            warn!("Firebase user cleanup not yet implemented");
            Ok(false)
        }
        Some(IdentityProvider::Custom) | None => {
            // For custom/unknown providers, try Keycloak if configured, else skip
            if config.keycloak_url.is_some() && config.keycloak_admin_username.is_some() {
                cleanup_keycloak(user, config, http).await
            } else {
                warn!(
                    username = user.username.as_deref(),
                    "No identity provider configured for cleanup — skipping"
                );
                Ok(false)
            }
        }
    }
}

/// Delete a user from Keycloak via the Admin REST API.
///
/// Flow: get admin token → search user by username → delete by ID.
async fn cleanup_keycloak(
    user: &TestUserRecord,
    config: &AgentConfig,
    http: &reqwest::Client,
) -> Result<bool, String> {
    let base_url = config
        .keycloak_url
        .as_deref()
        .ok_or("KEYCLOAK_URL not configured")?;
    let realm = config
        .keycloak_realm
        .as_deref()
        .ok_or("KEYCLOAK_REALM not configured")?;
    let admin_user = config
        .keycloak_admin_username
        .as_deref()
        .ok_or("KEYCLOAK_ADMIN_USERNAME not configured")?;
    let admin_pass = config
        .keycloak_admin_password
        .as_ref()
        .ok_or("KEYCLOAK_ADMIN_PASSWORD not configured")?;

    let username = user
        .username
        .as_deref()
        .ok_or("No username in test user record")?;

    info!(username, realm, "Cleaning up Keycloak test user");

    // Step 1: Get admin access token
    let token_url = format!("{base_url}/realms/master/protocol/openid-connect/token");
    let token_resp = http
        .post(&token_url)
        .form(&[
            ("grant_type", "password"),
            ("client_id", "admin-cli"),
            ("username", admin_user),
            ("password", admin_pass.expose_secret()),
        ])
        .send()
        .await
        .map_err(|e| format!("Keycloak token request failed: {e}"))?;

    if !token_resp.status().is_success() {
        let status = token_resp.status();
        let body = token_resp.text().await.unwrap_or_default();
        return Err(format!("Keycloak admin auth failed ({status}): {body}"));
    }

    let token_body: serde_json::Value = token_resp
        .json()
        .await
        .map_err(|e| format!("Failed to parse Keycloak token: {e}"))?;
    let access_token = token_body
        .get("access_token")
        .and_then(|v| v.as_str())
        .ok_or("No access_token in Keycloak response")?;

    // Step 2: Search for user by username
    let search_url =
        format!("{base_url}/admin/realms/{realm}/users?username={username}&exact=true");
    let search_resp = http
        .get(&search_url)
        .bearer_auth(access_token)
        .send()
        .await
        .map_err(|e| format!("Keycloak user search failed: {e}"))?;

    if !search_resp.status().is_success() {
        let status = search_resp.status();
        let body = search_resp.text().await.unwrap_or_default();
        return Err(format!("Keycloak user search failed ({status}): {body}"));
    }

    let users: Vec<serde_json::Value> = search_resp
        .json()
        .await
        .map_err(|e| format!("Failed to parse Keycloak users: {e}"))?;

    let user_id = users
        .first()
        .and_then(|u| u.get("id"))
        .and_then(|v| v.as_str())
        .ok_or_else(|| format!("User '{username}' not found in Keycloak realm '{realm}'"))?;

    // Step 3: Delete the user
    let delete_url = format!("{base_url}/admin/realms/{realm}/users/{user_id}");
    let delete_resp = http
        .delete(&delete_url)
        .bearer_auth(access_token)
        .send()
        .await
        .map_err(|e| format!("Keycloak user delete failed: {e}"))?;

    if delete_resp.status().is_success() || delete_resp.status().as_u16() == 204 {
        info!(username, user_id, "Keycloak test user deleted");
        Ok(true)
    } else {
        let status = delete_resp.status();
        let body = delete_resp.text().await.unwrap_or_default();
        Err(format!("Keycloak delete failed ({status}): {body}"))
    }
}

/// Delete a user from Auth0 via the Management API.
///
/// Requires `AUTH0_DOMAIN`, `AUTH0_CLIENT_ID`, `AUTH0_CLIENT_SECRET` env vars.
async fn cleanup_auth0(
    user: &TestUserRecord,
    _config: &AgentConfig,
    http: &reqwest::Client,
) -> Result<bool, String> {
    let domain = std::env::var("AUTH0_DOMAIN").map_err(|_| "AUTH0_DOMAIN not set")?;
    let client_id = std::env::var("AUTH0_CLIENT_ID").map_err(|_| "AUTH0_CLIENT_ID not set")?;
    let client_secret =
        std::env::var("AUTH0_CLIENT_SECRET").map_err(|_| "AUTH0_CLIENT_SECRET not set")?;

    let email = user
        .email
        .as_deref()
        .ok_or("No email in test user record for Auth0 lookup")?;

    info!(email, "Cleaning up Auth0 test user");

    // Get management API token
    let token_resp = http
        .post(format!("https://{domain}/oauth/token"))
        .json(&serde_json::json!({
            "grant_type": "client_credentials",
            "client_id": client_id,
            "client_secret": client_secret,
            "audience": format!("https://{domain}/api/v2/"),
        }))
        .send()
        .await
        .map_err(|e| format!("Auth0 token request failed: {e}"))?;

    let token_body: serde_json::Value = token_resp
        .json()
        .await
        .map_err(|e| format!("Failed to parse Auth0 token: {e}"))?;
    let access_token = token_body
        .get("access_token")
        .and_then(|v| v.as_str())
        .ok_or("No access_token in Auth0 response")?;

    // Search for user by email
    let encoded_email = urlencoding::encode(email);
    let search_url = format!("https://{domain}/api/v2/users-by-email?email={encoded_email}");
    let search_resp = http
        .get(&search_url)
        .bearer_auth(access_token)
        .send()
        .await
        .map_err(|e| format!("Auth0 user search failed: {e}"))?;

    let users: Vec<serde_json::Value> = search_resp
        .json()
        .await
        .map_err(|e| format!("Failed to parse Auth0 users: {e}"))?;

    let user_id = users
        .first()
        .and_then(|u| u.get("user_id"))
        .and_then(|v| v.as_str())
        .ok_or_else(|| format!("User with email '{email}' not found in Auth0"))?;

    // Delete
    let encoded_id = urlencoding::encode(user_id);
    let delete_resp = http
        .delete(format!("https://{domain}/api/v2/users/{encoded_id}"))
        .bearer_auth(access_token)
        .send()
        .await
        .map_err(|e| format!("Auth0 user delete failed: {e}"))?;

    if delete_resp.status().is_success() || delete_resp.status().as_u16() == 204 {
        info!(email, user_id, "Auth0 test user deleted");
        Ok(true)
    } else {
        let status = delete_resp.status();
        let body = delete_resp.text().await.unwrap_or_default();
        Err(format!("Auth0 delete failed ({status}): {body}"))
    }
}

/// Delete a user from Okta via the Users API.
///
/// Requires `OKTA_DOMAIN`, `OKTA_API_TOKEN` env vars.
async fn cleanup_okta(
    user: &TestUserRecord,
    _config: &AgentConfig,
    http: &reqwest::Client,
) -> Result<bool, String> {
    let domain = std::env::var("OKTA_DOMAIN").map_err(|_| "OKTA_DOMAIN not set")?;
    let api_token = std::env::var("OKTA_API_TOKEN").map_err(|_| "OKTA_API_TOKEN not set")?;

    let username = user
        .username
        .as_deref()
        .or(user.email.as_deref())
        .ok_or("No username/email in test user record for Okta lookup")?;

    info!(username, "Cleaning up Okta test user");

    // Search user
    let encoded = urlencoding::encode(username);
    let search_url = format!("https://{domain}/api/v1/users?search=profile.login+eq+\"{encoded}\"");
    let search_resp = http
        .get(&search_url)
        .header("Authorization", format!("SSWS {api_token}"))
        .send()
        .await
        .map_err(|e| format!("Okta user search failed: {e}"))?;

    let users: Vec<serde_json::Value> = search_resp
        .json()
        .await
        .map_err(|e| format!("Failed to parse Okta users: {e}"))?;

    let user_id = users
        .first()
        .and_then(|u| u.get("id"))
        .and_then(|v| v.as_str())
        .ok_or_else(|| format!("User '{username}' not found in Okta"))?;

    // Deactivate first (required by Okta before delete)
    let _ = http
        .post(format!(
            "https://{domain}/api/v1/users/{user_id}/lifecycle/deactivate"
        ))
        .header("Authorization", format!("SSWS {api_token}"))
        .send()
        .await;

    // Delete
    let delete_resp = http
        .delete(format!("https://{domain}/api/v1/users/{user_id}"))
        .header("Authorization", format!("SSWS {api_token}"))
        .send()
        .await
        .map_err(|e| format!("Okta user delete failed: {e}"))?;

    if delete_resp.status().is_success() || delete_resp.status().as_u16() == 204 {
        info!(username, user_id, "Okta test user deleted");
        Ok(true)
    } else {
        let status = delete_resp.status();
        let body = delete_resp.text().await.unwrap_or_default();
        Err(format!("Okta delete failed ({status}): {body}"))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use compliance_core::models::pentest::{IdentityProvider, TestUserRecord};
    use secrecy::SecretString;

    fn make_config_no_keycloak() -> AgentConfig {
        AgentConfig {
            mongodb_uri: String::new(),
            mongodb_database: String::new(),
            litellm_url: String::new(),
            litellm_api_key: SecretString::from(String::new()),
            litellm_model: String::new(),
            litellm_embed_model: String::new(),
            github_token: None,
            github_webhook_secret: None,
            gitlab_url: None,
            gitlab_token: None,
            gitlab_webhook_secret: None,
            jira_url: None,
            jira_email: None,
            jira_api_token: None,
            jira_project_key: None,
            searxng_url: None,
            nvd_api_key: None,
            agent_port: 3001,
            scan_schedule: String::new(),
            cve_monitor_schedule: String::new(),
            git_clone_base_path: String::new(),
            ssh_key_path: String::new(),
            keycloak_url: None,
            keycloak_realm: None,
            keycloak_admin_username: None,
            keycloak_admin_password: None,
            pentest_verification_email: None,
            pentest_imap_host: None,
            pentest_imap_port: None,
            pentest_imap_tls: true,
            pentest_imap_username: None,
            pentest_imap_password: None,
        }
    }

    #[tokio::test]
    async fn already_cleaned_up_returns_false() {
        let user = TestUserRecord {
            username: Some("test".into()),
            email: None,
            provider_user_id: None,
            provider: Some(IdentityProvider::Keycloak),
            cleaned_up: true,
        };
        let config = make_config_no_keycloak();
        let http = reqwest::Client::new();
        let result = cleanup_test_user(&user, &config, &http).await;
        assert_eq!(result, Ok(false));
    }

    #[tokio::test]
    async fn firebase_returns_false_not_implemented() {
        let user = TestUserRecord {
            username: Some("test".into()),
            email: None,
            provider_user_id: None,
            provider: Some(IdentityProvider::Firebase),
            cleaned_up: false,
        };
        let config = make_config_no_keycloak();
        let http = reqwest::Client::new();
        let result = cleanup_test_user(&user, &config, &http).await;
        assert_eq!(result, Ok(false));
    }

    #[tokio::test]
    async fn no_provider_no_keycloak_skips() {
        let user = TestUserRecord {
            username: Some("test".into()),
            email: None,
            provider_user_id: None,
            provider: None,
            cleaned_up: false,
        };
        let config = make_config_no_keycloak();
        let http = reqwest::Client::new();
        let result = cleanup_test_user(&user, &config, &http).await;
        assert_eq!(result, Ok(false));
    }

    #[tokio::test]
    async fn custom_provider_no_keycloak_skips() {
        let user = TestUserRecord {
            username: Some("test".into()),
            email: None,
            provider_user_id: None,
            provider: Some(IdentityProvider::Custom),
            cleaned_up: false,
        };
        let config = make_config_no_keycloak();
        let http = reqwest::Client::new();
        let result = cleanup_test_user(&user, &config, &http).await;
        assert_eq!(result, Ok(false));
    }

    #[tokio::test]
    async fn keycloak_missing_config_returns_error() {
        let user = TestUserRecord {
            username: Some("test".into()),
            email: None,
            provider_user_id: None,
            provider: Some(IdentityProvider::Keycloak),
            cleaned_up: false,
        };
        let config = make_config_no_keycloak();
        let http = reqwest::Client::new();
        let result = cleanup_test_user(&user, &config, &http).await;
        assert!(result.is_err());
        assert!(result
            .as_ref()
            .err()
            .is_some_and(|e| e.contains("KEYCLOAK_URL")));
    }

    #[tokio::test]
    async fn keycloak_missing_username_returns_error() {
        let user = TestUserRecord {
            username: None,
            email: Some("test@example.com".into()),
            provider_user_id: None,
            provider: Some(IdentityProvider::Keycloak),
            cleaned_up: false,
        };
        let mut config = make_config_no_keycloak();
        config.keycloak_url = Some("http://localhost:8080".into());
        config.keycloak_realm = Some("test".into());
        config.keycloak_admin_username = Some("admin".into());
        config.keycloak_admin_password = Some(SecretString::from("pass".to_string()));
        let http = reqwest::Client::new();
        let result = cleanup_test_user(&user, &config, &http).await;
        assert!(result.is_err());
        assert!(result
            .as_ref()
            .err()
            .is_some_and(|e| e.contains("username")));
    }

    #[tokio::test]
    async fn auth0_missing_env_returns_error() {
        let user = TestUserRecord {
            username: None,
            email: Some("test@example.com".into()),
            provider_user_id: None,
            provider: Some(IdentityProvider::Auth0),
            cleaned_up: false,
        };
        let config = make_config_no_keycloak();
        let http = reqwest::Client::new();
        let result = cleanup_test_user(&user, &config, &http).await;
        assert!(result.is_err());
        assert!(result
            .as_ref()
            .err()
            .is_some_and(|e| e.contains("AUTH0_DOMAIN")));
    }

    #[tokio::test]
    async fn okta_missing_env_returns_error() {
        let user = TestUserRecord {
            username: Some("test".into()),
            email: None,
            provider_user_id: None,
            provider: Some(IdentityProvider::Okta),
            cleaned_up: false,
        };
        let config = make_config_no_keycloak();
        let http = reqwest::Client::new();
        let result = cleanup_test_user(&user, &config, &http).await;
        assert!(result.is_err());
        assert!(result
            .as_ref()
            .err()
            .is_some_and(|e| e.contains("OKTA_DOMAIN")));
    }
}