exfiltrate 0.1.0

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Log capture and inspection tools for the transit proxy.
//!
//! This module provides functionality for capturing, storing, and searching through
//! logs generated by applications using the logwise logging framework. It enables
//! the proxy to intercept and store logs, making them available for inspection
//! through MCP tools.
//!
//! # Overview
//!
//! The module consists of:
//! - A centralized log storage system (`LogProxy`)
//! - Tools for reading logs (`LogwiseRead`)
//! - Tools for searching logs (`LogwiseGrep`)
//!
//! # Architecture
//!
//! Logs are stored in a thread-safe vector managed by `LogProxy`. The proxy
//! intercepts logwise output and stores it for later retrieval. Tools can then
//! query this storage to provide log inspection capabilities to MCP clients.
//!
//! # Log Format
//!
//! Captured logs typically follow the logwise format:
//! ```text
//! 0 INFO: examples/log_exfiltration.rs:4:5 [0ns] MESSAGE
//! ```
//!
//! Where:
//! - `0` is the task ID
//! - `INFO` is the log level (DEBUG, INFO, WARN, ERROR)
//! - `examples/log_exfiltration.rs:4:5` is the source location
//! - `[0ns]` is the timestamp
//! - `MESSAGE` is the actual log message
//!
//! # Usage
//!
//! The target application must call `exfiltrate::logwise::begin_capture()` to
//! start redirecting logs to this module. Logs generated before this call will
//! not be available for inspection.

use crate::mcp::tools::{Argument, InputSchema, Tool, ToolCallError, ToolCallResponse};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::{Arc, LazyLock, Mutex};

/// Global singleton instance of the log proxy.
///
/// This static instance is lazily initialized and provides the central
/// log storage for the entire application.
static CURRENT_LOGPROXY: LazyLock<LogProxy> = LazyLock::new(|| LogProxy::new());

/// Central log storage and management system.
///
/// `LogProxy` maintains a thread-safe collection of log messages captured
/// from the target application. It provides methods for adding new logs,
/// clearing the log buffer, and accessing the stored logs.
///
/// # Thread Safety
///
/// The internal storage uses `Arc<Mutex<Vec<String>>>` to ensure thread-safe
/// access from multiple concurrent contexts.
pub struct LogProxy {
    logs: Arc<Mutex<Vec<String>>>,
}

impl LogProxy {
    /// Returns the global singleton instance of the log proxy.
    ///
    /// This provides access to the centralized log storage system.
    pub fn current() -> &'static LogProxy {
        &CURRENT_LOGPROXY
    }

    /// Creates a new log proxy instance.
    ///
    /// This is called internally during static initialization.
    fn new() -> LogProxy {
        LogProxy {
            logs: Arc::new(Mutex::new(Vec::new())),
        }
    }

    /// Clears all stored logs.
    ///
    /// This removes all log entries from the internal storage, resetting
    /// the log buffer to an empty state.
    pub fn reset(&self) {
        self.logs.lock().unwrap().clear();
    }

    /// Adds a new log entry to the storage.
    ///
    /// Appends the provided log message to the end of the log buffer.
    /// Logs are stored in the order they are added.
    ///
    /// # Arguments
    ///
    /// * `log` - The log message to store
    pub fn add_log(&self, log: String) {
        self.logs.lock().unwrap().push(log);
    }
}

/// Response structure for the LogwiseRead tool.
///
/// Contains the requested log entries along with metadata about
/// the query and the total log count.
#[derive(Debug, serde::Serialize)]
struct LogResponse {
    /// The requested log entries
    logs: Vec<String>,
    /// Starting position in the log buffer (0-indexed)
    start_pos: usize,
    /// Ending position in the log buffer (exclusive)
    end_pos: usize,
    /// Total number of logs in storage
    all_logs: usize,
}
/// MCP tool for reading captured logs.
///
/// `LogwiseRead` provides access to logs stored by the proxy, allowing
/// clients to retrieve log entries with pagination support. This tool
/// is useful for debugging and monitoring applications through the proxy.
///
/// # Parameters
///
/// - `start_pos` (optional): Position to start reading from (0-indexed)
/// - `length` (optional): Number of logs to read (default: 10)
pub struct LogwiseRead;

impl Tool for LogwiseRead {
    fn name(&self) -> &str {
        "logwise_read"
    }

    fn description(&self) -> &str {
        "Reads logs from logwise.

        Often logs are printed to console.  However some environments with complex redirect setups
        may only print logs from certain threads or may not be flushed.  This tool allows
        more direct access to the logs.

        Limitations: in order for logs to be available from this tool, the target application must
        a) log with logwise, and b) call `exfiltrate::logwise::begin_capture()` to begin redirecting
        logs into this tool.  Logs made before this call will not be available.
        "
    }

    fn input_schema(&self) -> InputSchema {
        InputSchema::new(vec![
            Argument::new("start_pos".to_string(), "integer".to_string(), "The position to start reading logs from.  If omitted, tails the logs.".to_string(), false),
            Argument::new("length".to_string(), "integer".to_string(), "The number of logs to read.  If omitted, defaults to 10.  If the combination of start_pos and length go out of bounds, return as many logs are in bounds.".to_string(), false),
        ])
    }

    fn call(&self, params: HashMap<String, Value>) -> Result<ToolCallResponse, ToolCallError> {
        let log_proxy = LogProxy::current().logs.lock().unwrap();
        let length = params.get("length").and_then(|v| v.as_i64()).unwrap_or(10) as usize;

        let default_start_pos = log_proxy.len().saturating_sub(length);

        let start_pos = params
            .get("start_pos")
            .and_then(|v| v.as_i64())
            .map(|v| v as usize)
            .unwrap_or(default_start_pos);

        //adjust to make in bounds
        let start_pos = start_pos.min(log_proxy.len()).max(0);
        let end_pos = (start_pos + length).min(log_proxy.len());
        let logs = log_proxy[start_pos..end_pos].to_vec();
        let response = LogResponse {
            logs,
            start_pos,
            end_pos,
            all_logs: log_proxy.len(),
        };
        let response_text = serde_json::to_string(&response).unwrap();
        Ok(ToolCallResponse::new(vec![response_text.into()]))
    }
}
/// Represents a single log entry that matched a search pattern.
#[derive(Debug, serde::Serialize)]
struct MatchedLog {
    /// The full log message that matched
    log: String,
    /// Position of this log in the overall log buffer
    position: usize,
}

/// Response structure for the LogwiseGrep tool.
///
/// Contains all logs that matched the search pattern along with
/// their positions in the log buffer.
#[derive(Debug, serde::Serialize)]
struct LogwiseGrepResponse {
    /// Total number of logs in storage
    all_logs: usize,
    /// List of logs that matched the search pattern
    matched_logs: Vec<MatchedLog>,
}

/// MCP tool for searching through captured logs using regular expressions.
///
/// `LogwiseGrep` allows clients to search for specific patterns in the
/// stored logs using regular expressions. This is useful for filtering
/// logs by severity level, source location, or message content.
///
/// # Regular Expression Support
///
/// This tool uses the Rust `regex` crate, which supports:
/// - Basic patterns: `ERROR`, `WARN.*timeout`
/// - Character classes: `[0-9]+`, `\w+`
/// - Anchors: `^0 INFO`, `error$`
/// - Groups and alternation: `(ERROR|WARN)`, `task (\d+)`
///
pub struct LogwiseGrep;
impl Tool for LogwiseGrep {
    fn name(&self) -> &str {
        "logwise_grep"
    }

    fn description(&self) -> &str {
        "Greps logs from logwise.

        Limitations: same as logwise_read.
        "
    }

    fn input_schema(&self) -> InputSchema {
        let pattern_doc = r#"A regular expression to search for in logs.  If no logs match, returns an empty list.

        logwise_grep uses the `regex` crate for regular expressions, which supports a wide range of features.  For more information, see https://docs.rs/regex/latest/regex/

        Typical logwise logs are in the following format:
        ```
         0 INFO: examples/log_exfiltration.rs:4:5 [0ns] MESSAGE
        ```

        Where:
        - `0` is the task ID.  You can use this to filter logs by task ID.
        - `INFO` is the log level
        - `examples/log_exfiltration.rs:4:5` is the file and line
        - `[0ns]` is the timestamp
        - `MESSAGE` is the log message
        "#;

        InputSchema::new(vec![Argument::new(
            "pattern".to_string(),
            "string".to_string(),
            pattern_doc.to_string(),
            true,
        )])
    }

    fn call(&self, params: HashMap<String, Value>) -> Result<ToolCallResponse, ToolCallError> {
        let pattern = params
            .get("pattern")
            .and_then(|v| v.as_str())
            .ok_or_else(|| ToolCallError::new(vec!["No pattern".into()]))?;

        let regex = regex::Regex::new(pattern)
            .map_err(|_| ToolCallError::new(vec!["Invalid regex".into()]))?;
        let log_proxy = LogProxy::current().logs.lock().unwrap();

        let logs: Vec<MatchedLog> = log_proxy
            .iter()
            .enumerate()
            .filter_map(|(i, log)| {
                if regex.is_match(log) {
                    Some(MatchedLog {
                        log: log.clone(),
                        position: i,
                    })
                } else {
                    None
                }
            })
            .collect();

        let response = LogwiseGrepResponse {
            all_logs: log_proxy.len(),
            matched_logs: logs,
        };
        let res = serde_json::to_string(&response)
            .map_err(|e| ToolCallError::new(vec![e.to_string().into()]))?;
        Ok(ToolCallResponse::new(vec![res.into()]))
    }
}