compliance-scanner-agent/compliance-agent/src/main.rs

mod agent;
mod api;
pub(crate) 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>> {
    match dotenvy::dotenv() {
        Ok(path) => eprintln!("[dotenv] Loaded from: {}", path.display()),
        Err(e) => eprintln!("[dotenv] FAILED: {e}"),
    }

    let _telemetry_guard = compliance_core::telemetry::init_telemetry("compliance-agent");

    // Log critical env vars at startup
    tracing::info!(
        chrome_ws_url = std::env::var("CHROME_WS_URL").ok().as_deref(),
        pentest_email = std::env::var("PENTEST_VERIFICATION_EMAIL").ok().as_deref(),
        encryption_key_set = std::env::var("PENTEST_ENCRYPTION_KEY").is_ok(),
        "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(())
}