Files
compliance-scanner-agent/compliance-agent/src/main.rs
Sharang Parnerkar 6bc8ba89d1 feat: AI-driven automated penetration testing system
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>
2026-03-12 15:30:38 +01:00

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