111 lines
4.1 KiB
Rust
111 lines
4.1 KiB
Rust
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(())
|
|
}
|