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

use compliance_agent::{agent, api, config, database, migrate, scheduler, ssh, webhooks};

/// Run the `migrate onboarding` subcommand and exit. Backfills (or reverts) the
/// unified `onboarded_targets` collection per tenant.
///
/// Usage: `compliance-agent migrate onboarding [--all | --tenant <id>] [--dry-run] [--revert]`
async fn run_migration(
    args: &[String],
    pool: &database::DatabasePool,
) -> Result<(), compliance_agent::error::AgentError> {
    if args.get(2).map(String::as_str) != Some("onboarding") {
        eprintln!(
            "usage: compliance-agent migrate onboarding [--all | --tenant <id>] [--dry-run] [--revert]"
        );
        std::process::exit(2);
    }
    let has = |flag: &str| args.iter().any(|a| a == flag);
    let dry_run = has("--dry-run");
    let revert = has("--revert");
    let tenant = args
        .iter()
        .position(|a| a == "--tenant")
        .and_then(|i| args.get(i + 1))
        .cloned();

    let tenants: Vec<String> = if has("--all") {
        pool.list_tenant_ids().await?
    } else if let Some(t) = tenant {
        vec![t]
    } else {
        eprintln!("specify --all or --tenant <id>");
        std::process::exit(2);
    };

    for tenant_id in tenants {
        let db = pool.for_tenant_id(&tenant_id).await?;
        if revert {
            migrate::onboarding::revert(&db).await?;
            println!("[{tenant_id}] reverted onboarding backfill");
        } else {
            let report = migrate::onboarding::backfill_onboarded_targets(&db, dry_run).await?;
            let prefix = if dry_run { "(dry-run) " } else { "" };
            println!("[{tenant_id}] {prefix}{report:?}");
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    match dotenvy::dotenv() {
        Ok(path) => eprintln!("[dotenv] Loaded from: {}", path.display()),
        Err(_) => eprintln!("[dotenv] No .env file found, using environment variables"),
    }

    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...");
    // Per-tenant pool only — the agent has no shared "default" database
    // after M7.2-D. `mongodb_database` is now the db-name prefix used
    // for tenant databases (`<prefix>_<tenant_id>`).
    let db_pool =
        database::DatabasePool::connect(&config.mongodb_uri, &config.mongodb_database).await?;

    // One-shot subcommands run and exit without starting the servers.
    let args: Vec<String> = std::env::args().collect();
    if args.get(1).map(String::as_str) == Some("migrate") {
        run_migration(&args, &db_pool).await?;
        return Ok(());
    }

    let agent = agent::ComplianceAgent::new(config.clone(), db_pool);

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