compliance-scanner-agent/compliance-dast/src/tools/browser.rs

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use base64::Engine;
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.
///
/// **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 {
    fn default() -> Self {
        Self
    }
}

impl PentestTool for BrowserTool {
    fn name(&self) -> &str {
        "browser"
    }

    fn description(&self) -> &str {
        "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. 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 {
        json!({
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "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",
                    "description": "URL to navigate to (for 'navigate' action)"
                },
                "selector": {
                    "type": "string",
                    "description": "CSS selector for click/fill actions (e.g. '#username', 'a[href*=register]', 'button[type=submit]')"
                },
                "value": {
                    "type": "string",
                    "description": "Text value for 'fill' action, or JS expression for 'evaluate'"
                },
                "wait_ms": {
                    "type": "integer",
                    "description": "Milliseconds to wait after action (default: 1000)"
                }
            },
            "required": ["action"]
        })
    }

    fn execute<'a>(
        &'a self,
        input: serde_json::Value,
        context: &'a PentestToolContext,
    ) -> Pin<Box<dyn std::future::Future<Output = Result<PentestToolResult, CoreError>> + Send + 'a>>
    {
        Box::pin(async move {
            let action = input.get("action").and_then(|v| v.as_str()).unwrap_or("");
            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(1000);
            let session_key = context.session_id.clone();

            // 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,
                "screenshot" => session.screenshot().await,
                "click" => session.click(selector, wait_ms).await,
                "fill" => session.fill(selector, value, wait_ms).await,
                "get_content" => session.get_content().await,
                "evaluate" => session.evaluate(value).await,
                _ => Err(format!("Unknown browser action: {action}")),
            };

            // 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) => {
                    let summary = match action {
                        "navigate" => format!("Navigated to {url}"),
                        "screenshot" => "Captured page screenshot".to_string(),
                        "click" => format!("Clicked element: {selector}"),
                        "fill" => format!("Filled element: {selector}"),
                        "get_content" => "Retrieved page content".to_string(),
                        "evaluate" => "Evaluated JavaScript".to_string(),
                        _ => "Browser action completed".to_string(),
                    };
                    info!(action, %summary, "Browser tool executed");
                    Ok(PentestToolResult {
                        summary,
                        findings: Vec::new(),
                        data,
                    })
                }
                Err(e) => Ok(PentestToolResult {
                    summary: format!("Browser action '{action}' failed: {e}"),
                    findings: Vec::new(),
                    data: json!({ "error": e }),
                }),
            }
        })
    }
}

/// A single CDP session wrapping a browser tab.
struct BrowserSession {
    ws: WsStream,
    next_id: u64,
    session_id: String,
    target_id: String,
}

impl BrowserSession {
    /// Connect to headless Chrome and create a new tab.
    async fn connect() -> Result<Self, String> {
        let ws_url = std::env::var("CHROME_WS_URL").map_err(|_| {
            "CHROME_WS_URL not set — headless Chrome is required for browser actions".to_string()
        })?;

        // Discover browser WS endpoint
        let http_url = ws_url
            .replace("ws://", "http://")
            .replace("wss://", "https://");
        let version_url = format!("{http_url}/json/version");

        let version: serde_json::Value = reqwest::get(&version_url)
            .await
            .map_err(|e| format!("Cannot reach Chrome at {version_url}: {e}"))?
            .json()
            .await
            .map_err(|e| format!("Invalid /json/version response: {e}"))?;

        let browser_ws = version["webSocketDebuggerUrl"]
            .as_str()
            .ok_or_else(|| "No webSocketDebuggerUrl in /json/version".to_string())?;

        let (mut ws, _) = tokio_tungstenite::connect_async(browser_ws)
            .await
            .map_err(|e| format!("WebSocket connect failed: {e}"))?;

        let mut next_id: u64 = 1;

        // Create tab
        let resp = cdp_send(
            &mut ws,
            next_id,
            "Target.createTarget",
            json!({ "url": "about:blank" }),
        )
        .await?;
        next_id += 1;

        let target_id = resp
            .get("result")
            .and_then(|r| r.get("targetId"))
            .and_then(|t| t.as_str())
            .ok_or("No targetId in createTarget response")?
            .to_string();

        // Attach
        let resp = cdp_send(
            &mut ws,
            next_id,
            "Target.attachToTarget",
            json!({ "targetId": target_id, "flatten": true }),
        )
        .await?;
        next_id += 1;

        let session_id = resp
            .get("result")
            .and_then(|r| r.get("sessionId"))
            .and_then(|s| s.as_str())
            .ok_or("No sessionId in attachToTarget response")?
            .to_string();

        // Enable domains
        cdp_send_session(&mut ws, next_id, &session_id, "Page.enable", json!({})).await?;
        next_id += 1;

        cdp_send_session(&mut ws, next_id, &session_id, "Runtime.enable", json!({})).await?;
        next_id += 1;

        Ok(Self {
            ws,
            next_id,
            session_id,
            target_id,
        })
    }

    async fn navigate(&mut self, url: &str, wait_ms: u64) -> Result<serde_json::Value, String> {
        cdp_send_session(
            &mut self.ws,
            self.next_id,
            &self.session_id,
            "Page.navigate",
            json!({ "url": url }),
        )
        .await?;
        self.next_id += 1;

        tokio::time::sleep(Duration::from_millis(wait_ms)).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,
            "title": title,
            "elements": elements,
            "screenshot_base64": screenshot_b64,
        }))
    }

    /// 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,
            &self.session_id,
            "Page.captureScreenshot",
            json!({ "format": "png", "quality": 80 }),
        )
        .await?;
        self.next_id += 1;

        Ok(resp
            .get("result")
            .and_then(|r| r.get("data"))
            .and_then(|d| d.as_str())
            .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)
            .map(|b| b.len() / 1024)
            .unwrap_or(0);

        Ok(json!({
            "screenshot_base64": b64,
            "size_kb": size_kb,
        }))
    }

    async fn click(&mut self, selector: &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}"}});
                var rect = el.getBoundingClientRect();
                el.click();
                return JSON.stringify({{
                    clicked: true,
                    tag: el.tagName,
                    text: el.textContent.substring(0, 100),
                    x: rect.x + rect.width/2,
                    y: rect.y + rect.height/2
                }});
            }})()"#,
            sel = serde_json::to_string(selector).unwrap_or_default(),
            raw = selector.replace('"', r#"\""#),
        );

        let result = self.evaluate_raw(&js).await?;
        tokio::time::sleep(Duration::from_millis(wait_ms)).await;

        // 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(
        &mut self,
        selector: &str,
        value: &str,
        wait_ms: u64,
    ) -> Result<serde_json::Value, String> {
        // 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(),
        );
        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();

        tokio::time::sleep(Duration::from_millis(wait_ms)).await;

        Ok(json!({
            "filled": true,
            "selector": selector,
            "value": final_value,
        }))
    }

    async fn get_content(&mut self) -> Result<serde_json::Value, String> {
        let title = self
            .evaluate_raw("document.title")
            .await
            .unwrap_or_default();
        let url = self
            .evaluate_raw("window.location.href")
            .await
            .unwrap_or_default();

        // 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,
            "page": page_data,
        }))
    }

    async fn evaluate(&mut self, expression: &str) -> Result<serde_json::Value, String> {
        let result = self.evaluate_raw(expression).await?;
        Ok(json!({
            "result": result,
        }))
    }

    /// Execute JS and return the string result.
    async fn evaluate_raw(&mut self, expression: &str) -> Result<String, String> {
        let resp = cdp_send_session(
            &mut self.ws,
            self.next_id,
            &self.session_id,
            "Runtime.evaluate",
            json!({
                "expression": expression,
                "returnByValue": true,
            }),
        )
        .await?;
        self.next_id += 1;

        let result = resp
            .get("result")
            .and_then(|r| r.get("result"))
            .and_then(|r| r.get("value"));

        match result {
            Some(serde_json::Value::String(s)) => Ok(s.clone()),
            Some(v) => Ok(v.to_string()),
            None => Ok(String::new()),
        }
    }

    async fn close(&mut self) -> Result<(), String> {
        let _ = cdp_send(
            &mut self.ws,
            self.next_id,
            "Target.closeTarget",
            json!({ "targetId": self.target_id }),
        )
        .await;
        let _ = self.ws.close(None).await;
        Ok(())
    }
}

/// 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,
    id: u64,
    method: &str,
    params: serde_json::Value,
) -> Result<serde_json::Value, String> {
    let msg = json!({ "id": id, "method": method, "params": params });
    ws.send(Message::Text(msg.to_string().into()))
        .await
        .map_err(|e| format!("WS send failed: {e}"))?;
    read_until_result(ws, id).await
}

async fn cdp_send_session(
    ws: &mut WsStream,
    id: u64,
    session_id: &str,
    method: &str,
    params: serde_json::Value,
) -> Result<serde_json::Value, String> {
    let msg = json!({
        "id": id,
        "sessionId": session_id,
        "method": method,
        "params": params,
    });
    ws.send(Message::Text(msg.to_string().into()))
        .await
        .map_err(|e| format!("WS send failed: {e}"))?;
    read_until_result(ws, id).await
}

async fn read_until_result(ws: &mut WsStream, id: u64) -> Result<serde_json::Value, String> {
    let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
    loop {
        let msg = tokio::time::timeout_at(deadline, ws.next())
            .await
            .map_err(|_| format!("Timeout waiting for CDP response id={id}"))?
            .ok_or_else(|| "WebSocket closed unexpectedly".to_string())?
            .map_err(|e| format!("WebSocket read error: {e}"))?;

        if let Message::Text(text) = msg {
            if let Ok(val) = serde_json::from_str::<serde_json::Value>(&text) {
                if val.get("id").and_then(|i| i.as_u64()) == Some(id) {
                    if let Some(err) = val.get("error") {
                        return Err(format!("CDP error: {err}"));
                    }
                    return Ok(val);
                }
            }
        }
    }
}