compliance-scanner-agent/compliance-agent/src/api/handlers/onboarding.rs

//! Onboarding API — CRUD for unified targets, artifact add, classification, and
//! the scan-applicability matrix. The wizard (and future integrations) drive
//! onboarding through these endpoints.

use std::collections::HashMap;
use std::sync::Arc;

use axum::extract::{Extension, Multipart, Path, Query};
use axum::http::StatusCode;
use axum::Json;
use mongodb::bson::{doc, oid::ObjectId, to_bson};
use serde::{Deserialize, Serialize};

use compliance_core::models::{
    Artifact, ArtifactKind, ComplianceProfile, OnboardedTarget, PlcFormat, TargetScanConfig,
    TargetType,
};
use compliance_core::scan_matrix::{applicable_scans, supports_pentest};
use compliance_core::tenant_ctx::TenantCtx;

use crate::agent::ComplianceAgent;
use crate::classify::{classify_target, MockFirmwareDetector};

use super::dto::tenant_db;
use super::{collect_cursor_async, ApiResponse, PaginationParams};

type AgentExt = Extension<Arc<ComplianceAgent>>;

/// A client-supplied artifact spec. The server builds the [`Artifact`] (and its
/// id) from it, so clients never set internal fields.
#[derive(Deserialize)]
pub struct ArtifactInput {
    pub kind: ArtifactKind,
    pub source_ref: String,
    #[serde(default)]
    pub branch: Option<String>,
    #[serde(default)]
    pub plc_format: Option<PlcFormat>,
}

impl ArtifactInput {
    fn build(&self) -> Artifact {
        let s = self.source_ref.clone();
        match self.kind {
            ArtifactKind::GitRepo => {
                Artifact::git_repo(s, self.branch.clone().unwrap_or_else(|| "main".to_string()))
            }
            ArtifactKind::LiveUrl => Artifact::live_url(s),
            ArtifactKind::FirmwareImage => Artifact::firmware_image(s),
            ArtifactKind::SourceArchive => Artifact::source_archive(s),
            ArtifactKind::MobilePackage => Artifact::mobile_package(s),
            ArtifactKind::ContainerImage => Artifact::container_image(s),
            ArtifactKind::PlcProject => {
                Artifact::plc_project(s, self.plc_format.unwrap_or(PlcFormat::PlcopenXml))
            }
            ArtifactKind::PlaintextDescription => Artifact::plaintext(s),
        }
    }
}

#[derive(Deserialize)]
pub struct CreateTargetRequest {
    pub name: String,
    pub target_type: TargetType,
    #[serde(default)]
    pub description: Option<String>,
    #[serde(default)]
    pub artifacts: Vec<ArtifactInput>,
}

#[derive(Deserialize)]
pub struct UpdateTargetRequest {
    pub name: Option<String>,
    pub target_type: Option<TargetType>,
    pub scan_config: Option<TargetScanConfig>,
    pub compliance_profile: Option<ComplianceProfile>,
    pub scan_schedule: Option<String>,
    /// Replace the target's artifacts wholesale (used by the dashboard editor).
    #[serde(default)]
    pub artifacts: Option<Vec<ArtifactInput>>,
}

/// One applicable-scan option, serialized for the wizard.
#[derive(Serialize)]
pub struct ScanOptionDto {
    pub scan: String,
    pub default_on: bool,
    pub rationale: String,
    pub required_artifact: Option<String>,
    pub blocked_reason: Option<String>,
}

#[derive(Serialize)]
pub struct ApplicableScansResponse {
    pub scans: Vec<ScanOptionDto>,
    pub pentest_supported: bool,
}

fn parse_oid(id: &str) -> Result<ObjectId, StatusCode> {
    ObjectId::parse_str(id).map_err(|_| StatusCode::BAD_REQUEST)
}

/// GET /api/v1/targets — list onboarded targets (paginated).
#[tracing::instrument(skip_all)]
pub async fn list_targets(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Query(params): Query<PaginationParams>,
) -> Result<Json<ApiResponse<Vec<OnboardedTarget>>>, StatusCode> {
    let db = tenant_db(&agent, &tenant).await?;
    let skip = (params.page.saturating_sub(1)) * params.limit as u64;
    let total = db
        .onboarded_targets()
        .count_documents(doc! {})
        .await
        .unwrap_or(0);
    let targets = match db
        .onboarded_targets()
        .find(doc! {})
        .skip(skip)
        .limit(params.limit)
        .await
    {
        Ok(cursor) => collect_cursor_async(cursor).await,
        Err(e) => {
            tracing::warn!("Failed to fetch onboarded targets: {e}");
            Vec::new()
        }
    };
    Ok(Json(ApiResponse {
        data: targets,
        total: Some(total),
        page: Some(params.page),
    }))
}

/// POST /api/v1/targets — create an onboarded target.
#[tracing::instrument(skip_all)]
pub async fn create_target(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Json(req): Json<CreateTargetRequest>,
) -> Result<Json<ApiResponse<OnboardedTarget>>, StatusCode> {
    let mut target = OnboardedTarget::new(req.name, req.target_type);
    target.description = req.description;
    target.artifacts = req.artifacts.iter().map(ArtifactInput::build).collect();

    let db = tenant_db(&agent, &tenant).await?;
    let res = db
        .onboarded_targets()
        .insert_one(&target)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    target.id = res.inserted_id.as_object_id();
    Ok(Json(ApiResponse {
        data: target,
        total: None,
        page: None,
    }))
}

/// GET /api/v1/targets/{id} — fetch one target.
#[tracing::instrument(skip_all, fields(target_id = %id))]
pub async fn get_target(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
) -> Result<Json<ApiResponse<OnboardedTarget>>, StatusCode> {
    let oid = parse_oid(&id)?;
    let db = tenant_db(&agent, &tenant).await?;
    let target = db
        .onboarded_targets()
        .find_one(doc! { "_id": oid })
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;
    Ok(Json(ApiResponse {
        data: target,
        total: None,
        page: None,
    }))
}

/// PATCH /api/v1/targets/{id} — update mutable fields.
#[tracing::instrument(skip_all, fields(target_id = %id))]
pub async fn update_target(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
    Json(req): Json<UpdateTargetRequest>,
) -> Result<Json<ApiResponse<OnboardedTarget>>, StatusCode> {
    let oid = parse_oid(&id)?;
    let db = tenant_db(&agent, &tenant).await?;

    let mut set = doc! { "updated_at": mongodb::bson::DateTime::now() };
    if let Some(name) = req.name {
        set.insert("name", name);
    }
    if let Some(tt) = req.target_type {
        set.insert(
            "target_type",
            to_bson(&tt).map_err(|_| StatusCode::BAD_REQUEST)?,
        );
    }
    if let Some(sc) = req.scan_config {
        set.insert(
            "scan_config",
            to_bson(&sc).map_err(|_| StatusCode::BAD_REQUEST)?,
        );
    }
    if let Some(cp) = req.compliance_profile {
        set.insert(
            "compliance_profile",
            to_bson(&cp).map_err(|_| StatusCode::BAD_REQUEST)?,
        );
    }
    if let Some(ss) = req.scan_schedule {
        set.insert("scan_schedule", ss);
    }
    if let Some(arts) = req.artifacts {
        let built: Vec<Artifact> = arts.iter().map(ArtifactInput::build).collect();
        set.insert(
            "artifacts",
            to_bson(&built).map_err(|_| StatusCode::BAD_REQUEST)?,
        );
    }

    db.onboarded_targets()
        .update_one(doc! { "_id": oid }, doc! { "$set": set })
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    get_target(Extension(agent), tenant, Path(id)).await
}

/// DELETE /api/v1/targets/{id} — remove the target and its findings/scans.
#[tracing::instrument(skip_all, fields(target_id = %id))]
pub async fn delete_target(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, StatusCode> {
    let oid = parse_oid(&id)?;
    let db = tenant_db(&agent, &tenant).await?;
    db.onboarded_targets()
        .delete_one(doc! { "_id": oid })
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    // Cascade all data keyed by repo_id == target id (best-effort).
    let db = &db;
    let _ = db.findings().delete_many(doc! { "repo_id": &id }).await;
    let _ = db.sbom_entries().delete_many(doc! { "repo_id": &id }).await;
    let _ = db.scan_runs().delete_many(doc! { "repo_id": &id }).await;
    let _ = db.cve_alerts().delete_many(doc! { "repo_id": &id }).await;
    let _ = db
        .tracker_issues()
        .delete_many(doc! { "repo_id": &id })
        .await;
    let _ = db.graph_nodes().delete_many(doc! { "repo_id": &id }).await;
    let _ = db.graph_edges().delete_many(doc! { "repo_id": &id }).await;
    let _ = db.graph_builds().delete_many(doc! { "repo_id": &id }).await;
    let _ = db
        .impact_analyses()
        .delete_many(doc! { "repo_id": &id })
        .await;
    let _ = db
        .code_embeddings()
        .delete_many(doc! { "repo_id": &id })
        .await;
    let _ = db
        .embedding_builds()
        .delete_many(doc! { "repo_id": &id })
        .await;

    // DAST targets linked to this target, and all their downstream data.
    if let Ok(mut cursor) = db.dast_targets().find(doc! { "repo_id": &id }).await {
        use futures_util::StreamExt;
        while let Some(Ok(dt)) = cursor.next().await {
            let dast_target_id = dt.id.map(|oid| oid.to_hex()).unwrap_or_default();
            if !dast_target_id.is_empty() {
                cascade_delete_dast_target(db, &dast_target_id).await;
            }
        }
    }

    // Pentest sessions linked directly to this target (not via a DAST target).
    if let Ok(mut cursor) = db.pentest_sessions().find(doc! { "repo_id": &id }).await {
        use futures_util::StreamExt;
        while let Some(Ok(session)) = cursor.next().await {
            let session_id = session.id.map(|oid| oid.to_hex()).unwrap_or_default();
            if !session_id.is_empty() {
                let _ = db
                    .attack_chain_nodes()
                    .delete_many(doc! { "session_id": &session_id })
                    .await;
                let _ = db
                    .pentest_messages()
                    .delete_many(doc! { "session_id": &session_id })
                    .await;
                let _ = db
                    .dast_findings()
                    .delete_many(doc! { "session_id": &session_id })
                    .await;
            }
        }
    }
    let _ = db
        .pentest_sessions()
        .delete_many(doc! { "repo_id": &id })
        .await;

    Ok(Json(serde_json::json!({ "status": "deleted" })))
}

/// Delete a DAST target and everything downstream of it (pentest sessions +
/// their attack chains / messages / findings, DAST scan runs + findings).
async fn cascade_delete_dast_target(db: &crate::database::Database, target_id: &str) {
    use futures_util::StreamExt;
    if let Ok(mut cursor) = db
        .pentest_sessions()
        .find(doc! { "target_id": target_id })
        .await
    {
        while let Some(Ok(session)) = cursor.next().await {
            let session_id = session.id.map(|oid| oid.to_hex()).unwrap_or_default();
            if !session_id.is_empty() {
                let _ = db
                    .attack_chain_nodes()
                    .delete_many(doc! { "session_id": &session_id })
                    .await;
                let _ = db
                    .pentest_messages()
                    .delete_many(doc! { "session_id": &session_id })
                    .await;
                let _ = db
                    .dast_findings()
                    .delete_many(doc! { "session_id": &session_id })
                    .await;
            }
        }
    }
    let _ = db
        .pentest_sessions()
        .delete_many(doc! { "target_id": target_id })
        .await;
    let _ = db
        .dast_findings()
        .delete_many(doc! { "target_id": target_id })
        .await;
    let _ = db
        .dast_scan_runs()
        .delete_many(doc! { "target_id": target_id })
        .await;
    if let Ok(oid) = mongodb::bson::oid::ObjectId::parse_str(target_id) {
        let _ = db.dast_targets().delete_one(doc! { "_id": oid }).await;
    }
}

/// POST /api/v1/targets/{id}/artifacts — attach an artifact (by reference).
#[tracing::instrument(skip_all, fields(target_id = %id))]
pub async fn add_artifact(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
    Json(input): Json<ArtifactInput>,
) -> Result<Json<ApiResponse<OnboardedTarget>>, StatusCode> {
    let oid = parse_oid(&id)?;
    let db = tenant_db(&agent, &tenant).await?;
    let artifact = to_bson(&input.build()).map_err(|_| StatusCode::BAD_REQUEST)?;
    db.onboarded_targets()
        .update_one(
            doc! { "_id": oid },
            doc! { "$push": { "artifacts": artifact }, "$set": { "updated_at": mongodb::bson::DateTime::now() } },
        )
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    get_target(Extension(agent), tenant, Path(id)).await
}

/// POST /api/v1/targets/{id}/artifacts/upload — attach an artifact by uploading
/// its file (PLC project, firmware image, source archive, mobile package). The
/// bytes are written to the artifact blob store and referenced by `stored_path`,
/// so ingest resolves them locally (no URL fetch).
///
/// Multipart fields: `file` (required), `kind` (required, snake_case
/// `ArtifactKind`), `plc_format` (optional, for PLC projects).
#[tracing::instrument(skip_all, fields(target_id = %id))]
pub async fn upload_artifact(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
    mut multipart: Multipart,
) -> Result<Json<ApiResponse<OnboardedTarget>>, StatusCode> {
    let oid = parse_oid(&id)?;
    let db = tenant_db(&agent, &tenant).await?;
    if db
        .onboarded_targets()
        .find_one(doc! { "_id": oid })
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .is_none()
    {
        return Err(StatusCode::NOT_FOUND);
    }

    let mut kind: Option<ArtifactKind> = None;
    let mut plc_format: Option<PlcFormat> = None;
    let mut filename = String::from("upload.bin");
    let mut bytes: Option<axum::body::Bytes> = None;

    while let Some(field) = multipart
        .next_field()
        .await
        .map_err(|_| StatusCode::BAD_REQUEST)?
    {
        match field.name().unwrap_or("") {
            "kind" => {
                let v = field.text().await.map_err(|_| StatusCode::BAD_REQUEST)?;
                kind = parse_enum(&v);
            }
            "plc_format" => {
                let v = field.text().await.map_err(|_| StatusCode::BAD_REQUEST)?;
                plc_format = parse_enum(&v);
            }
            "file" => {
                if let Some(fname) = field.file_name() {
                    filename = fname.to_string();
                }
                bytes = Some(field.bytes().await.map_err(|_| StatusCode::BAD_REQUEST)?);
            }
            _ => {}
        }
    }

    let (Some(kind), Some(bytes)) = (kind, bytes) else {
        return Err(StatusCode::BAD_REQUEST);
    };

    // Store the uploaded bytes under the artifact blob store.
    let safe_name: String = filename
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | '_') {
                c
            } else {
                '_'
            }
        })
        .collect();
    let dir = std::path::Path::new(&agent.config.artifact_store_base_path)
        .join("uploads")
        .join(&id);
    std::fs::create_dir_all(&dir).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    let dest = dir.join(format!("{}_{safe_name}", uuid::Uuid::new_v4()));
    std::fs::write(&dest, bytes.as_ref()).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    // Build the artifact for this kind, referencing the stored file.
    let mut artifact = match kind {
        ArtifactKind::PlcProject => Artifact::plc_project(
            filename.clone(),
            plc_format.unwrap_or(PlcFormat::PlcopenXml),
        ),
        ArtifactKind::FirmwareImage => Artifact::firmware_image(filename.clone()),
        ArtifactKind::SourceArchive => Artifact::source_archive(filename.clone()),
        ArtifactKind::MobilePackage => Artifact::mobile_package(filename.clone()),
        // Non-file kinds (git repo, live URL, container ref, text) use the JSON
        // add-artifact endpoint, not upload.
        _ => return Err(StatusCode::BAD_REQUEST),
    };
    artifact.stored_path = Some(dest.to_string_lossy().to_string());
    artifact.size_bytes = Some(bytes.len() as u64);

    let artifact_bson = to_bson(&artifact).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    db.onboarded_targets()
        .update_one(
            doc! { "_id": oid },
            doc! { "$push": { "artifacts": artifact_bson }, "$set": { "updated_at": mongodb::bson::DateTime::now() } },
        )
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    get_target(Extension(agent), tenant, Path(id)).await
}

/// Deserialize a snake_case enum value from a plain string.
fn parse_enum<T: for<'de> Deserialize<'de>>(s: &str) -> Option<T> {
    serde_json::from_value(serde_json::Value::String(s.to_string())).ok()
}

/// GET /api/v1/targets/{id}/applicable-scans — the scan-applicability matrix.
#[tracing::instrument(skip_all, fields(target_id = %id))]
pub async fn applicable_scans_for_target(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
) -> Result<Json<ApiResponse<ApplicableScansResponse>>, StatusCode> {
    let oid = parse_oid(&id)?;
    let db = tenant_db(&agent, &tenant).await?;
    let target = db
        .onboarded_targets()
        .find_one(doc! { "_id": oid })
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;

    let scans = applicable_scans(&target)
        .into_iter()
        .map(|o| ScanOptionDto {
            scan: o.scan.to_string(),
            default_on: o.default_on,
            rationale: o.rationale,
            required_artifact: o.required_artifact.map(|k| k.to_string()),
            blocked_reason: o.blocked_reason,
        })
        .collect();

    Ok(Json(ApiResponse {
        data: ApplicableScansResponse {
            scans,
            pentest_supported: supports_pentest(target.target_type),
        },
        total: None,
        page: None,
    }))
}

/// POST /api/v1/targets/{id}/detect — classify the target from its artifacts.
///
/// This is the lightweight pass: it classifies from artifact kinds without
/// ingesting (cloning) sources, so it returns immediately. Deep detection (after
/// ingest, with tramiton firmware analysis) is a follow-up background step.
#[tracing::instrument(skip_all, fields(target_id = %id))]
pub async fn detect_target(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
) -> Result<Json<ApiResponse<OnboardedTarget>>, StatusCode> {
    let oid = parse_oid(&id)?;
    let db = tenant_db(&agent, &tenant).await?;
    let mut target = db
        .onboarded_targets()
        .find_one(doc! { "_id": oid })
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;

    // No ingested working paths here → kind-based classification only; the mock
    // firmware detector is never invoked (no firmware working path present).
    let empty = HashMap::new();
    let detector = MockFirmwareDetector { detection: None };
    let classification = classify_target(&target, &empty, &detector)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let classification_bson =
        to_bson(&classification).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    db.onboarded_targets()
        .update_one(
            doc! { "_id": oid },
            doc! { "$set": { "classification": classification_bson, "updated_at": mongodb::bson::DateTime::now() } },
        )
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    target.classification = Some(classification);

    Ok(Json(ApiResponse {
        data: target,
        total: None,
        page: None,
    }))
}

/// POST /api/v1/targets/{id}/scan — trigger a scan for the target.
///
/// Dispatches to the unified pipeline when `UNIFIED_PIPELINE` is set (else the
/// legacy path). Runs in the background and returns immediately.
#[tracing::instrument(skip_all, fields(target_id = %id))]
pub async fn trigger_target_scan(
    Extension(agent): AgentExt,
    tenant: TenantCtx,
    Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, StatusCode> {
    let oid = parse_oid(&id)?;
    let db = tenant_db(&agent, &tenant).await?;
    // 404 if the target doesn't exist for this tenant.
    if db
        .onboarded_targets()
        .find_one(doc! { "_id": oid })
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .is_none()
    {
        return Err(StatusCode::NOT_FOUND);
    }

    let agent_clone = (*agent).clone();
    let tenant_id = tenant.0.tenant_id.clone();
    tokio::spawn(async move {
        // Always the unified target pipeline — this endpoint is about an
        // onboarded target by construction, independent of the global
        // `unified_pipeline` transition flag used by the legacy paths.
        if let Err(e) = agent_clone
            .run_target_scan(
                &tenant_id,
                &id,
                compliance_core::models::ScanTrigger::Manual,
            )
            .await
        {
            tracing::error!("Manual target scan failed for {id}: {e}");
        }
    });
    Ok(Json(serde_json::json!({ "status": "scan_triggered" })))
}