Add a complete AI pentest system where Claude autonomously drives security testing via tool-calling. The LLM selects from 16 tools, chains results, and builds an attack chain DAG. Core: - PentestTool trait (dyn-compatible) with PentestToolContext/Result - PentestSession, AttackChainNode, PentestMessage, PentestEvent models - 10 new DastVulnType variants (DNS, DMARC, TLS, cookies, CSP, CORS, etc.) - LLM client chat_with_tools() for OpenAI-compatible tool calling Tools (16 total): - 5 agent wrappers: SQL injection, XSS, auth bypass, SSRF, API fuzzer - 11 new infra tools: DNS checker, DMARC checker, TLS analyzer, security headers, cookie analyzer, CSP analyzer, rate limit tester, console log detector, CORS checker, OpenAPI parser, recon - ToolRegistry for tool lookup and LLM definition generation Orchestrator: - PentestOrchestrator with iterative tool-calling loop (max 50 rounds) - Attack chain node recording per tool invocation - SSE event broadcasting for real-time progress - Strategy-aware system prompts (quick/comprehensive/targeted/aggressive/stealth) API (9 endpoints): - POST/GET /pentest/sessions, GET /pentest/sessions/:id - POST /pentest/sessions/:id/chat, GET /pentest/sessions/:id/stream - GET /pentest/sessions/:id/attack-chain, messages, findings - GET /pentest/stats Dashboard: - Pentest dashboard with stat cards, severity distribution, session list - Chat-based session page with split layout (chat + findings/attack chain) - Inline tool execution indicators, auto-polling, new session modal - Sidebar navigation item Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
59 lines
1.7 KiB
Rust
59 lines
1.7 KiB
Rust
mod agent;
|
|
mod api;
|
|
mod config;
|
|
mod database;
|
|
mod error;
|
|
mod llm;
|
|
mod pentest;
|
|
mod pipeline;
|
|
mod rag;
|
|
mod scheduler;
|
|
mod ssh;
|
|
#[allow(dead_code)]
|
|
mod trackers;
|
|
mod webhooks;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
dotenvy::dotenv().ok();
|
|
|
|
let _telemetry_guard = compliance_core::telemetry::init_telemetry("compliance-agent");
|
|
|
|
tracing::info!("Loading configuration...");
|
|
let config = config::load_config()?;
|
|
|
|
// Ensure SSH key pair exists for cloning private repos
|
|
match ssh::ensure_ssh_key(&config.ssh_key_path) {
|
|
Ok(pubkey) => tracing::info!("SSH public key: {}", pubkey.trim()),
|
|
Err(e) => tracing::warn!("SSH key generation skipped: {e}"),
|
|
}
|
|
|
|
tracing::info!("Connecting to MongoDB...");
|
|
let db = database::Database::connect(&config.mongodb_uri, &config.mongodb_database).await?;
|
|
db.ensure_indexes().await?;
|
|
|
|
let agent = agent::ComplianceAgent::new(config.clone(), db.clone());
|
|
|
|
tracing::info!("Starting scheduler...");
|
|
let scheduler_agent = agent.clone();
|
|
let scheduler_handle = tokio::spawn(async move {
|
|
if let Err(e) = scheduler::start_scheduler(&scheduler_agent).await {
|
|
tracing::error!("Scheduler error: {e}");
|
|
}
|
|
});
|
|
|
|
tracing::info!("Starting webhook server...");
|
|
let webhook_agent = agent.clone();
|
|
let webhook_handle = tokio::spawn(async move {
|
|
if let Err(e) = webhooks::start_webhook_server(&webhook_agent).await {
|
|
tracing::error!("Webhook server error: {e}");
|
|
}
|
|
});
|
|
|
|
tracing::info!("Starting REST API on port {}...", config.agent_port);
|
|
api::start_api_server(agent, config.agent_port).await?;
|
|
|
|
let _ = tokio::join!(scheduler_handle, webhook_handle);
|
|
Ok(())
|
|
}
|