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

pub mod api_fuzzer;
pub mod auth_bypass;
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)));

        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()
    }
}