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

pub mod api_fuzzer;
pub mod auth_bypass;
pub mod browser;
pub mod console_log_detector;
pub mod cookie_analyzer;
pub mod cors_checker;
pub mod csp_analyzer;
pub mod dmarc_checker;
pub mod dns_checker;
pub mod openapi_parser;
pub mod rate_limit_tester;
pub mod recon;
pub mod security_headers;
pub mod sql_injection;
pub mod ssrf;
pub mod tls_analyzer;
pub mod xss;

use std::collections::HashMap;

use compliance_core::traits::pentest_tool::PentestTool;

/// A definition describing a tool for LLM tool_use registration.
#[derive(Debug, Clone)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub input_schema: serde_json::Value,
}

/// Registry that holds all available pentest tools and provides
/// look-up by name.
pub struct ToolRegistry {
    tools: HashMap<String, Box<dyn PentestTool>>,
}

impl Default for ToolRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl ToolRegistry {
    /// Create a new registry with all built-in tools pre-registered.
    #[allow(clippy::expect_used)]
    pub fn new() -> Self {
        let http = reqwest::Client::builder()
            .danger_accept_invalid_certs(true)
            .timeout(std::time::Duration::from_secs(30))
            .redirect(reqwest::redirect::Policy::limited(5))
            .build()
            .expect("failed to build HTTP client");

        let mut tools: HashMap<String, Box<dyn PentestTool>> = HashMap::new();

        // Agent-wrapping tools
        let register = |tools: &mut HashMap<String, Box<dyn PentestTool>>,
                        tool: Box<dyn PentestTool>| {
            tools.insert(tool.name().to_string(), tool);
        };

        register(
            &mut tools,
            Box::new(sql_injection::SqlInjectionTool::new(http.clone())),
        );
        register(&mut tools, Box::new(xss::XssTool::new(http.clone())));
        register(
            &mut tools,
            Box::new(auth_bypass::AuthBypassTool::new(http.clone())),
        );
        register(&mut tools, Box::new(ssrf::SsrfTool::new(http.clone())));
        register(
            &mut tools,
            Box::new(api_fuzzer::ApiFuzzerTool::new(http.clone())),
        );

        // New infrastructure / analysis tools
        register(&mut tools, Box::<dns_checker::DnsCheckerTool>::default());
        register(
            &mut tools,
            Box::<dmarc_checker::DmarcCheckerTool>::default(),
        );
        register(
            &mut tools,
            Box::new(tls_analyzer::TlsAnalyzerTool::new(http.clone())),
        );
        register(
            &mut tools,
            Box::new(security_headers::SecurityHeadersTool::new(http.clone())),
        );
        register(
            &mut tools,
            Box::new(cookie_analyzer::CookieAnalyzerTool::new(http.clone())),
        );
        register(
            &mut tools,
            Box::new(csp_analyzer::CspAnalyzerTool::new(http.clone())),
        );
        register(
            &mut tools,
            Box::new(rate_limit_tester::RateLimitTesterTool::new(http.clone())),
        );
        register(
            &mut tools,
            Box::new(console_log_detector::ConsoleLogDetectorTool::new(
                http.clone(),
            )),
        );
        register(
            &mut tools,
            Box::new(cors_checker::CorsCheckerTool::new(http.clone())),
        );
        register(
            &mut tools,
            Box::new(openapi_parser::OpenApiParserTool::new(http.clone())),
        );
        register(&mut tools, Box::new(recon::ReconTool::new(http)));
        register(&mut tools, Box::<browser::BrowserTool>::default());

        Self { tools }
    }

    /// Look up a tool by name.
    pub fn get(&self, name: &str) -> Option<&dyn PentestTool> {
        self.tools.get(name).map(|b| b.as_ref())
    }

    /// Return definitions for every registered tool.
    pub fn all_definitions(&self) -> Vec<ToolDefinition> {
        self.tools
            .values()
            .map(|t| ToolDefinition {
                name: t.name().to_string(),
                description: t.description().to_string(),
                input_schema: t.input_schema(),
            })
            .collect()
    }

    /// Return the names of all registered tools.
    pub fn list_names(&self) -> Vec<String> {
        self.tools.keys().cloned().collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn registry_has_all_expected_tools() {
        let registry = ToolRegistry::new();
        let names = registry.list_names();

        let expected = [
            "recon",
            "openapi_parser",
            "dns_checker",
            "dmarc_checker",
            "tls_analyzer",
            "security_headers",
            "cookie_analyzer",
            "csp_analyzer",
            "cors_checker",
            "rate_limit_tester",
            "console_log_detector",
            "sql_injection_scanner",
            "xss_scanner",
            "ssrf_scanner",
            "auth_bypass_scanner",
            "api_fuzzer",
            "browser",
        ];

        for name in &expected {
            assert!(
                names.contains(&name.to_string()),
                "Missing tool: {name}. Registered: {names:?}"
            );
        }
    }

    #[test]
    fn registry_get_returns_tool() {
        let registry = ToolRegistry::new();
        assert!(registry.get("recon").is_some());
        assert!(registry.get("browser").is_some());
        assert!(registry.get("nonexistent").is_none());
    }

    #[test]
    fn all_definitions_have_valid_schemas() {
        let registry = ToolRegistry::new();
        let defs = registry.all_definitions();

        assert!(!defs.is_empty());
        for def in &defs {
            assert!(!def.name.is_empty(), "Tool has empty name");
            assert!(
                !def.description.is_empty(),
                "Tool {} has empty description",
                def.name
            );
            assert!(
                def.input_schema.is_object(),
                "Tool {} schema is not an object",
                def.name
            );
            // Every schema should have "type": "object"
            assert_eq!(
                def.input_schema.get("type").and_then(|v| v.as_str()),
                Some("object"),
                "Tool {} schema type is not 'object'",
                def.name
            );
        }
    }

    #[test]
    fn browser_tool_schema_has_action_enum() {
        let registry = ToolRegistry::new();
        let browser = registry.get("browser");
        assert!(browser.is_some());
        let schema = browser.map(|t| t.input_schema()).unwrap_or_default();
        let action_prop = schema.get("properties").and_then(|p| p.get("action"));
        assert!(
            action_prop.is_some(),
            "Browser tool missing 'action' property"
        );
        let action_enum = action_prop
            .and_then(|a| a.get("enum"))
            .and_then(|e| e.as_array());
        assert!(action_enum.is_some(), "Browser action missing enum");
        let actions: Vec<&str> = action_enum
            .into_iter()
            .flatten()
            .filter_map(|v| v.as_str())
            .collect();
        assert!(actions.contains(&"navigate"));
        assert!(actions.contains(&"screenshot"));
        assert!(actions.contains(&"click"));
        assert!(actions.contains(&"fill"));
        assert!(actions.contains(&"get_content"));
        assert!(actions.contains(&"close"));
    }
}