firefox-webdriver 0.1.4

//! Network interception and blocking methods.

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use base64::Engine;
use base64::engine::general_purpose::STANDARD as Base64Standard;
use serde_json::Value;
use tracing::debug;

use crate::browser::network::{
    BodyAction, HeadersAction, InterceptedRequest, InterceptedRequestBody,
    InterceptedRequestHeaders, InterceptedResponse, InterceptedResponseBody, RequestAction,
    RequestBody, ResponseBodyData,
};
use crate::error::{Error, Result};
use crate::identifiers::InterceptId;
use crate::protocol::{Command, Event, EventReply, NetworkCommand, Response};

use super::Tab;

// ============================================================================
// Constants
// ============================================================================

const EVENT_BEFORE_REQUEST: &str = "network.beforeRequestSent";
const EVENT_REQUEST_HEADERS: &str = "network.requestHeaders";
const EVENT_REQUEST_BODY: &str = "network.requestBody";
const EVENT_RESPONSE_HEADERS: &str = "network.responseHeaders";
const EVENT_RESPONSE_BODY: &str = "network.responseBody";

/// Global counter for generating unique handler keys per intercept registration.
static HANDLER_KEY_COUNTER: AtomicU64 = AtomicU64::new(0);

/// Generates a unique handler key for an intercept registration.
fn next_handler_key(prefix: &str) -> String {
    let id = HANDLER_KEY_COUNTER.fetch_add(1, Ordering::Relaxed);
    format!("{prefix}_{id}")
}

// ============================================================================
// InterceptBuilder
// ============================================================================

/// Builder for configuring network intercept filters.
///
/// # Example
///
/// ```ignore
/// let id = tab.intercept()
///     .url_patterns(vec!["*api*".into()])
///     .resource_types(vec!["xhr".into(), "fetch".into()])
///     .requests(|req| {
///         RequestAction::allow()
///     }).await?;
/// ```
pub struct InterceptBuilder<'a> {
    tab: &'a Tab,
    url_patterns: Option<Vec<String>>,
    resource_types: Option<Vec<String>>,
}

impl<'a> InterceptBuilder<'a> {
    /// Sets URL patterns to filter intercepted requests/responses.
    ///
    /// Patterns support glob-style wildcards (`*`).
    #[inline]
    #[must_use]
    pub fn url_patterns(mut self, patterns: Vec<String>) -> Self {
        self.url_patterns = Some(patterns);
        self
    }

    /// Sets resource types to filter intercepted requests/responses.
    ///
    /// Valid types: "document", "script", "stylesheet", "image", "media",
    /// "font", "xhr", "fetch", "websocket", "other".
    #[inline]
    #[must_use]
    pub fn resource_types(mut self, types: Vec<String>) -> Self {
        self.resource_types = Some(types);
        self
    }

    /// Intercepts requests with the configured filters.
    ///
    /// # Returns
    ///
    /// An `InterceptId` that can be used to stop this intercept.
    pub async fn requests<F>(self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedRequest) -> RequestAction + Send + Sync + 'static,
    {
        debug!(tab_id = %self.tab.inner.tab_id, "Enabling filtered request interception");

        let window = self.tab.get_window()?;
        let callback = Arc::new(callback);
        let handler_key = next_handler_key("intercept_request");

        window.inner.pool.add_event_handler(
            window.inner.session_id,
            handler_key.clone(),
            Box::new(move |event: Event| {
                if event.method.as_str() != EVENT_BEFORE_REQUEST {
                    return None;
                }

                let request = parse_intercepted_request(&event);
                let action = callback(request);
                let result = request_action_to_json(&action);

                Some(EventReply::new(
                    event.id,
                    EVENT_BEFORE_REQUEST,
                    result,
                ))
            }),
        );

        let command = Command::Network(NetworkCommand::AddIntercept {
            intercept_requests: true,
            intercept_request_headers: false,
            intercept_request_body: false,
            intercept_responses: false,
            intercept_response_body: false,
            url_patterns: self.url_patterns,
            resource_types: self.resource_types,
        });

        let response = self.tab.send_command(command).await?;
        let intercept_id = extract_intercept_id(&response)?;
        window.inner.intercept_handlers.lock().insert(intercept_id.clone(), handler_key);
        Ok(intercept_id)
    }

    /// Intercepts request headers with the configured filters.
    pub async fn request_headers<F>(self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedRequestHeaders) -> HeadersAction + Send + Sync + 'static,
    {
        debug!(tab_id = %self.tab.inner.tab_id, "Enabling filtered request headers interception");

        let window = self.tab.get_window()?;
        let callback = Arc::new(callback);
        let handler_key = next_handler_key("intercept_request_headers");

        window.inner.pool.add_event_handler(
            window.inner.session_id,
            handler_key.clone(),
            Box::new(move |event: Event| {
                if event.method.as_str() != EVENT_REQUEST_HEADERS {
                    return None;
                }

                let headers_data = parse_intercepted_request_headers(&event);
                let action = callback(headers_data);
                let result = headers_action_to_json(&action);

                Some(EventReply::new(event.id, EVENT_REQUEST_HEADERS, result))
            }),
        );

        let command = Command::Network(NetworkCommand::AddIntercept {
            intercept_requests: false,
            intercept_request_headers: true,
            intercept_request_body: false,
            intercept_responses: false,
            intercept_response_body: false,
            url_patterns: self.url_patterns,
            resource_types: self.resource_types,
        });

        let response = self.tab.send_command(command).await?;
        let intercept_id = extract_intercept_id(&response)?;
        window.inner.intercept_handlers.lock().insert(intercept_id.clone(), handler_key);
        Ok(intercept_id)
    }

    /// Intercepts request body with the configured filters (read-only).
    pub async fn request_body<F>(self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedRequestBody) + Send + Sync + 'static,
    {
        debug!(tab_id = %self.tab.inner.tab_id, "Enabling filtered request body interception");

        let window = self.tab.get_window()?;
        let callback = Arc::new(callback);
        let handler_key = next_handler_key("intercept_request_body");

        window.inner.pool.add_event_handler(
            window.inner.session_id,
            handler_key.clone(),
            Box::new(move |event: Event| {
                if event.method.as_str() != EVENT_REQUEST_BODY {
                    return None;
                }

                let body_data = parse_intercepted_request_body(&event);
                callback(body_data);

                Some(EventReply::new(
                    event.id,
                    EVENT_REQUEST_BODY,
                    serde_json::json!({ "action": "allow" }),
                ))
            }),
        );

        let command = Command::Network(NetworkCommand::AddIntercept {
            intercept_requests: false,
            intercept_request_headers: false,
            intercept_request_body: true,
            intercept_responses: false,
            intercept_response_body: false,
            url_patterns: self.url_patterns,
            resource_types: self.resource_types,
        });

        let response = self.tab.send_command(command).await?;
        let intercept_id = extract_intercept_id(&response)?;
        window.inner.intercept_handlers.lock().insert(intercept_id.clone(), handler_key);
        Ok(intercept_id)
    }

    /// Intercepts response headers with the configured filters.
    pub async fn responses<F>(self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedResponse) -> HeadersAction + Send + Sync + 'static,
    {
        debug!(tab_id = %self.tab.inner.tab_id, "Enabling filtered response interception");

        let window = self.tab.get_window()?;
        let callback = Arc::new(callback);
        let handler_key = next_handler_key("intercept_response");

        window.inner.pool.add_event_handler(
            window.inner.session_id,
            handler_key.clone(),
            Box::new(move |event: Event| {
                if event.method.as_str() != EVENT_RESPONSE_HEADERS {
                    return None;
                }

                let resp = parse_intercepted_response(&event);
                let action = callback(resp);
                let result = headers_action_to_json(&action);

                Some(EventReply::new(event.id, EVENT_RESPONSE_HEADERS, result))
            }),
        );

        let command = Command::Network(NetworkCommand::AddIntercept {
            intercept_requests: false,
            intercept_request_headers: false,
            intercept_request_body: false,
            intercept_responses: true,
            intercept_response_body: false,
            url_patterns: self.url_patterns,
            resource_types: self.resource_types,
        });

        let response = self.tab.send_command(command).await?;
        let intercept_id = extract_intercept_id(&response)?;
        window.inner.intercept_handlers.lock().insert(intercept_id.clone(), handler_key);
        Ok(intercept_id)
    }

    /// Intercepts response body with the configured filters.
    pub async fn response_body<F>(self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedResponseBody) -> BodyAction + Send + Sync + 'static,
    {
        debug!(tab_id = %self.tab.inner.tab_id, "Enabling filtered response body interception");

        let window = self.tab.get_window()?;
        let callback = Arc::new(callback);
        let handler_key = next_handler_key("intercept_response_body");

        window.inner.pool.add_event_handler(
            window.inner.session_id,
            handler_key.clone(),
            Box::new(move |event: Event| {
                if event.method.as_str() != EVENT_RESPONSE_BODY {
                    return None;
                }

                let body_data = parse_intercepted_response_body(&event);
                let action = callback(body_data);
                let result = body_action_to_json(&action);

                Some(EventReply::new(event.id, EVENT_RESPONSE_BODY, result))
            }),
        );

        let command = Command::Network(NetworkCommand::AddIntercept {
            intercept_requests: false,
            intercept_request_headers: false,
            intercept_request_body: false,
            intercept_responses: false,
            intercept_response_body: true,
            url_patterns: self.url_patterns,
            resource_types: self.resource_types,
        });

        let response = self.tab.send_command(command).await?;
        let intercept_id = extract_intercept_id(&response)?;
        window.inner.intercept_handlers.lock().insert(intercept_id.clone(), handler_key);
        Ok(intercept_id)
    }
}

// ============================================================================
// Tab - Network
// ============================================================================

impl Tab {
    /// Creates an intercept builder for filtered network interception.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let id = tab.intercept()
    ///     .url_patterns(vec!["*api*".into()])
    ///     .resource_types(vec!["xhr".into(), "fetch".into()])
    ///     .requests(|req| RequestAction::allow())
    ///     .await?;
    /// ```
    #[must_use]
    pub fn intercept(&self) -> InterceptBuilder<'_> {
        InterceptBuilder {
            tab: self,
            url_patterns: None,
            resource_types: None,
        }
    }

    /// Sets URL patterns to block.
    ///
    /// Patterns support wildcards (`*`).
    ///
    /// # Example
    ///
    /// ```ignore
    /// tab.set_block_rules(&["*ads*", "*tracking*"]).await?;
    /// ```
    pub async fn set_block_rules(&self, patterns: &[&str]) -> Result<()> {
        debug!(tab_id = %self.inner.tab_id, pattern_count = patterns.len(), "Setting block rules");

        let command = Command::Network(NetworkCommand::SetBlockRules {
            patterns: patterns.iter().map(|s| (*s).to_string()).collect(),
        });

        self.send_command(command).await?;
        Ok(())
    }

    /// Clears all URL block rules.
    pub async fn clear_block_rules(&self) -> Result<()> {
        debug!(tab_id = %self.inner.tab_id, "Clearing block rules");
        let command = Command::Network(NetworkCommand::ClearBlockRules);
        self.send_command(command).await?;
        Ok(())
    }

    /// Intercepts network requests with a callback.
    ///
    /// # Returns
    ///
    /// An `InterceptId` that can be used to stop this intercept.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use firefox_webdriver::RequestAction;
    ///
    /// let id = tab.intercept_request(|req| {
    ///     if req.url.contains("ads") {
    ///         RequestAction::block()
    ///     } else {
    ///         RequestAction::allow()
    ///     }
    /// }).await?;
    /// ```
    pub async fn intercept_request<F>(&self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedRequest) -> RequestAction + Send + Sync + 'static,
    {
        self.intercept().requests(callback).await
    }

    /// Intercepts request headers with a callback.
    pub async fn intercept_request_headers<F>(&self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedRequestHeaders) -> HeadersAction + Send + Sync + 'static,
    {
        self.intercept().request_headers(callback).await
    }

    /// Intercepts request body for logging (read-only).
    pub async fn intercept_request_body<F>(&self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedRequestBody) + Send + Sync + 'static,
    {
        self.intercept().request_body(callback).await
    }

    /// Intercepts response headers with a callback.
    pub async fn intercept_response<F>(&self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedResponse) -> HeadersAction + Send + Sync + 'static,
    {
        self.intercept().responses(callback).await
    }

    /// Intercepts response body with a callback.
    pub async fn intercept_response_body<F>(&self, callback: F) -> Result<InterceptId>
    where
        F: Fn(InterceptedResponseBody) -> BodyAction + Send + Sync + 'static,
    {
        self.intercept().response_body(callback).await
    }

    /// Stops network interception.
    ///
    /// Only removes the specific handler associated with this intercept,
    /// leaving other intercept handlers intact.
    ///
    /// # Arguments
    ///
    /// * `intercept_id` - The intercept ID returned from intercept methods
    pub async fn stop_intercept(&self, intercept_id: &InterceptId) -> Result<()> {
        debug!(tab_id = %self.inner.tab_id, %intercept_id, "Stopping interception");

        let window = self.get_window()?;

        if let Some(handler_key) = window.inner.intercept_handlers.lock().remove(intercept_id) {
            window
                .inner
                .pool
                .remove_event_handler(window.inner.session_id, &handler_key);
        }

        let command = Command::Network(NetworkCommand::RemoveIntercept {
            intercept_id: intercept_id.clone(),
        });

        self.send_command(command).await?;
        Ok(())
    }
}

// ============================================================================
// JSON Field Extraction Helpers
// ============================================================================

/// Extracts a string field from event params.
fn get_str(params: &serde_json::Value, key: &str) -> String {
    params.get(key).and_then(|v| v.as_str()).unwrap_or("").to_string()
}

/// Extracts a string field with a default value.
fn get_str_or(params: &serde_json::Value, key: &str, default: &str) -> String {
    params.get(key).and_then(|v| v.as_str()).unwrap_or(default).to_string()
}

/// Extracts a u32 field from event params.
fn get_u32(params: &serde_json::Value, key: &str) -> u32 {
    params.get(key).and_then(|v| v.as_u64()).unwrap_or(0) as u32
}

/// Extracts a u64 field from event params.
fn get_u64(params: &serde_json::Value, key: &str) -> u64 {
    params.get(key).and_then(|v| v.as_u64()).unwrap_or(0)
}

/// Extracts a u16 field from event params.
fn get_u16(params: &serde_json::Value, key: &str) -> u16 {
    params.get(key).and_then(|v| v.as_u64()).unwrap_or(0) as u16
}

/// Parses headers from event params.
fn parse_headers(params: &serde_json::Value) -> HashMap<String, String> {
    params.get("headers")
        .and_then(|v| v.as_object())
        .map(|obj| {
            obj.iter()
                .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
                .collect()
        })
        .unwrap_or_default()
}

// ============================================================================
// Helper Functions
// ============================================================================

/// Extracts intercept ID from response.
fn extract_intercept_id(response: &Response) -> Result<InterceptId> {
    let id = response
        .result
        .as_ref()
        .and_then(|v| v.get("interceptId"))
        .and_then(|v| v.as_str())
        .ok_or_else(|| Error::protocol("No interceptId in response"))?;

    Ok(InterceptId::new(id))
}

/// Parses intercepted request from event.
fn parse_intercepted_request(event: &Event) -> InterceptedRequest {
    InterceptedRequest {
        request_id: get_str(&event.params, "requestId"),
        url: get_str(&event.params, "url"),
        method: get_str_or(&event.params, "method", "GET"),
        headers: parse_headers(&event.params),
        resource_type: get_str_or(&event.params, "resourceType", "other"),
        tab_id: get_u32(&event.params, "tabId"),
        frame_id: get_u64(&event.params, "frameId"),
        body: None,
    }
}

/// Parses intercepted request headers from event.
fn parse_intercepted_request_headers(event: &Event) -> InterceptedRequestHeaders {
    InterceptedRequestHeaders {
        request_id: get_str(&event.params, "requestId"),
        url: get_str(&event.params, "url"),
        method: get_str_or(&event.params, "method", "GET"),
        headers: parse_headers(&event.params),
        tab_id: get_u32(&event.params, "tabId"),
        frame_id: get_u64(&event.params, "frameId"),
    }
}

/// Parses intercepted request body from event.
fn parse_intercepted_request_body(event: &Event) -> InterceptedRequestBody {
    InterceptedRequestBody {
        request_id: get_str(&event.params, "requestId"),
        url: get_str(&event.params, "url"),
        method: get_str_or(&event.params, "method", "GET"),
        resource_type: get_str_or(&event.params, "resourceType", "other"),
        tab_id: get_u32(&event.params, "tabId"),
        frame_id: get_u64(&event.params, "frameId"),
        body: event.params.as_object().and_then(parse_request_body),
    }
}

/// Parses intercepted response from event.
fn parse_intercepted_response(event: &Event) -> InterceptedResponse {
    InterceptedResponse {
        request_id: get_str(&event.params, "requestId"),
        url: get_str(&event.params, "url"),
        status: get_u16(&event.params, "status"),
        status_text: get_str(&event.params, "statusText"),
        headers: parse_headers(&event.params),
        tab_id: get_u32(&event.params, "tabId"),
        frame_id: get_u64(&event.params, "frameId"),
    }
}

/// Parses intercepted response body from event.
fn parse_intercepted_response_body(event: &Event) -> InterceptedResponseBody {
    let body = if let Some(b64) = event.params.get("bodyBase64").and_then(|v| v.as_str()) {
        match Base64Standard.decode(b64) {
            Ok(bytes) => ResponseBodyData::Binary(bytes),
            Err(_) => ResponseBodyData::Text(String::new()),
        }
    } else {
        let text = get_str(&event.params, "body");
        ResponseBodyData::Text(text)
    };

    InterceptedResponseBody {
        request_id: get_str(&event.params, "requestId"),
        url: get_str(&event.params, "url"),
        status: event
            .params
            .get("status")
            .and_then(|v| v.as_u64())
            .unwrap_or(200) as u16,
        content_type: get_str_or(&event.params, "contentType", "application/octet-stream"),
        body,
        tab_id: get_u32(&event.params, "tabId"),
        frame_id: get_u64(&event.params, "frameId"),
        content_length: get_u64(&event.params, "contentLength") as usize,
    }
}

/// Parses request body from event params.
fn parse_request_body(params: &serde_json::Map<String, Value>) -> Option<RequestBody> {
    let body = params.get("body")?;
    let body_obj = body.as_object()?;

    if let Some(error) = body_obj.get("error").and_then(|v| v.as_str()) {
        return Some(RequestBody::Error(error.to_string()));
    }

    if let Some(form_data) = body_obj.get("data").and_then(|v| v.as_object())
        && body_obj.get("type").and_then(|v| v.as_str()) == Some("formData")
    {
        let mut map = HashMap::new();
        for (key, value) in form_data {
            if let Some(arr) = value.as_array() {
                let values: Vec<String> = arr
                    .iter()
                    .filter_map(|v| v.as_str().map(|s| s.to_string()))
                    .collect();
                map.insert(key.clone(), values);
            }
        }
        return Some(RequestBody::FormData(map));
    }

    if let Some(raw_data) = body_obj.get("data").and_then(|v| v.as_array())
        && body_obj.get("type").and_then(|v| v.as_str()) == Some("raw")
    {
        let mut bytes = Vec::new();
        for item in raw_data {
            if let Some(obj) = item.as_object()
                && let Some(b64) = obj.get("data").and_then(|v| v.as_str())
                && let Ok(decoded) = Base64Standard.decode(b64)
            {
                bytes.extend(decoded);
            }
        }
        if !bytes.is_empty() {
            return Some(RequestBody::Raw(bytes));
        }
    }

    None
}

/// Converts request action to JSON.
fn request_action_to_json(action: &RequestAction) -> Value {
    match action {
        RequestAction::Allow => serde_json::json!({ "action": "allow" }),
        RequestAction::Block => serde_json::json!({ "action": "block" }),
        RequestAction::Redirect(url) => serde_json::json!({ "action": "redirect", "url": url }),
    }
}

/// Converts headers action to JSON.
fn headers_action_to_json(action: &HeadersAction) -> Value {
    match action {
        HeadersAction::Allow => serde_json::json!({ "action": "allow" }),
        HeadersAction::Block => serde_json::json!({ "action": "block" }),
        HeadersAction::Modify {
            headers,
            status_code,
        } => {
            let mut json = serde_json::json!({ "action": "modifyHeaders", "headers": headers });
            if let Some(code) = status_code {
                json["statusCode"] = serde_json::json!(code);
            }
            json
        }
    }
}

/// Converts body action to JSON.
fn body_action_to_json(action: &BodyAction) -> Value {
    match action {
        BodyAction::Allow => serde_json::json!({ "action": "allow" }),
        BodyAction::ModifyBody(b) => serde_json::json!({ "action": "modifyBody", "body": b }),
    }
}