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>, } 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> = HashMap::new(); // Agent-wrapping tools let register = |tools: &mut HashMap>, tool: Box| { 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::::default()); register( &mut tools, Box::::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 { 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 { self.tools.keys().cloned().collect() } }