charybdis/src/lib.rs

use std::sync::Arc;
use std::time::Instant;
use tonic::{Request, Response, Status};
use tracing::{error, info, instrument};
use uuid::Uuid;

pub mod adapters;
pub mod config;
pub mod database;
pub mod error;
pub mod events;
pub mod findings;
pub mod plugins;
pub mod scanners;
pub mod security;
pub mod telemetry;

// Include the generated protobuf code with serde support
pub mod charybdis {
    pub mod core {
        tonic::include_proto!("charybdis.core");
    }
    pub mod entities {
        tonic::include_proto!("charybdis.entities");
    }
    pub mod ingestion {
        tonic::include_proto!("charybdis.ingestion");
    }
    pub mod plugins {
        pub mod defectdojo {
            tonic::include_proto!("charybdis.plugins.defectdojo");
        }
        pub mod dependencytrack {
            tonic::include_proto!("charybdis.plugins.dependencytrack");
        }
        pub mod keycloak {
            tonic::include_proto!("charybdis.plugins.keycloak");
        }
    }
}

use charybdis::entities::entity_service_server::EntityService;

// Export the compiled protobuf descriptor set for reflection
// This is used by binaries that need to set up gRPC reflection
pub static ENTITY_DESCRIPTOR_SET: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/entity_descriptor.bin"));
use charybdis::entities::{
    CreateEntityRequest, CreateEntityResponse, DeleteEntityRequest, DeleteEntityResponse,
    GetEntityRequest, GetEntityResponse, ListEntitiesRequest, ListEntitiesResponse,
    UpdateEntityRequest, UpdateEntityResponse,
};

use database::EntityRepository;
use events::{EntityEvent, EventBus};
use security::interceptor::AuthInterceptor;
use telemetry::Metrics;

#[derive(Clone)]
pub struct MyEntityService {
    repository: EntityRepository,
    event_bus: Arc<dyn EventBus>,
    metrics: Metrics,
    auth_interceptor: Option<AuthInterceptor>,
}

impl MyEntityService {
    pub fn new(
        repository: EntityRepository,
        event_bus: Arc<dyn EventBus>,
        metrics: Metrics,
        auth_interceptor: Option<AuthInterceptor>,
    ) -> Self {
        Self {
            repository,
            event_bus,
            metrics,
            auth_interceptor,
        }
    }
}

#[tonic::async_trait]
impl EntityService for MyEntityService {
    #[instrument(skip(self, request), fields(entity.kind))]
    async fn create_entity(
        &self,
        request: Request<CreateEntityRequest>,
    ) -> Result<Response<CreateEntityResponse>, Status> {
        let start = Instant::now();

        let request = if let Some(ref interceptor) = self.auth_interceptor {
            interceptor
                .authorize_request(request, "/charybdis.entities.EntityService/CreateEntity")?
        } else {
            request
        };

        let entity_data = request
            .into_inner()
            .entity
            .ok_or_else(|| Status::invalid_argument("Entity data is missing"))?;

        validate_entity(&entity_data)?;

        let entity_kind = entity_data.kind.clone();
        tracing::Span::current().record("entity.kind", &entity_kind);

        let created_entity = self
            .repository
            .create(&entity_data)
            .await
            .map_err(|e| {
                error!(error = %e, "Failed to create entity");
                Status::from(e)
            })?;

        let entity_id = created_entity.id.clone();
        info!(entity_id = %entity_id, entity_kind = %entity_kind, "Entity created");

        if let Ok(uuid) = Uuid::parse_str(&entity_id) {
            let event = EntityEvent::created(uuid)
                .with_metadata("entity_kind".to_string(), created_entity.kind.clone())
                .with_entity_data(Arc::new(created_entity.clone()));

            if let Err(e) = self.event_bus.publish(event).await {
                error!(entity_id = %entity_id, error = %e, "Failed to publish entity created event");
            }
        }

        let duration = start.elapsed().as_secs_f64();
        self.metrics
            .record_entity_operation("create", &entity_kind, duration);

        Ok(Response::new(CreateEntityResponse {
            entity: Some(created_entity),
        }))
    }

    #[instrument(skip(self, request), fields(entity.id))]
    async fn get_entity(
        &self,
        request: Request<GetEntityRequest>,
    ) -> Result<Response<GetEntityResponse>, Status> {
        let start = Instant::now();

        let request = if let Some(ref interceptor) = self.auth_interceptor {
            interceptor.authorize_request(request, "/charybdis.entities.EntityService/GetEntity")?
        } else {
            request
        };

        let entity_id = request.into_inner().id;
        tracing::Span::current().record("entity.id", &entity_id);

        let entity = self
            .repository
            .get_by_id(&entity_id)
            .await
            .map_err(|e| {
                error!(entity_id = %entity_id, error = %e, "Failed to get entity");
                Status::from(e)
            })?
            .ok_or_else(|| Status::not_found(format!("Entity not found: {}", entity_id)))?;

        let duration = start.elapsed().as_secs_f64();
        self.metrics
            .record_entity_operation("get", &entity.kind, duration);

        Ok(Response::new(GetEntityResponse {
            entity: Some(entity),
        }))
    }

    #[instrument(skip(self, request), fields(entity.id, entity.kind))]
    async fn update_entity(
        &self,
        request: Request<UpdateEntityRequest>,
    ) -> Result<Response<UpdateEntityResponse>, Status> {
        let start = Instant::now();

        let request = if let Some(ref interceptor) = self.auth_interceptor {
            interceptor
                .authorize_request(request, "/charybdis.entities.EntityService/UpdateEntity")?
        } else {
            request
        };

        let req = request.into_inner();
        let entity_id = req.id.clone();

        let updated_entity_data = req
            .entity
            .ok_or_else(|| Status::invalid_argument("Entity data is missing"))?;

        validate_entity(&updated_entity_data)?;

        let entity_kind = updated_entity_data.kind.clone();
        tracing::Span::current().record("entity.id", &entity_id);
        tracing::Span::current().record("entity.kind", &entity_kind);

        let update_result = if let Some(field_mask) = req.update_mask {
            self.repository
                .partial_update(&entity_id, &updated_entity_data, &field_mask)
                .await
        } else {
            self.repository
                .update(&entity_id, &updated_entity_data)
                .await
        };

        let updated_entity = update_result
            .map_err(|e| {
                error!(entity_id = %entity_id, error = %e, "Failed to update entity");
                Status::from(e)
            })?
            .ok_or_else(|| {
                Status::not_found(format!("Entity not found: {}", entity_id))
            })?;

        info!(entity_id = %entity_id, entity_kind = %entity_kind, "Entity updated");

        if let Ok(uuid) = Uuid::parse_str(&entity_id) {
            let event = EntityEvent::updated(uuid)
                .with_metadata("entity_kind".to_string(), updated_entity.kind.clone())
                .with_entity_data(Arc::new(updated_entity.clone()));

            if let Err(e) = self.event_bus.publish(event).await {
                error!(entity_id = %entity_id, error = %e, "Failed to publish entity updated event");
            }
        }

        let duration = start.elapsed().as_secs_f64();
        self.metrics
            .record_entity_operation("update", &entity_kind, duration);

        Ok(Response::new(UpdateEntityResponse {
            entity: Some(updated_entity),
        }))
    }

    #[instrument(skip(self, request), fields(entity.id, entity.kind))]
    async fn delete_entity(
        &self,
        request: Request<DeleteEntityRequest>,
    ) -> Result<Response<DeleteEntityResponse>, Status> {
        let start = Instant::now();

        let request = if let Some(ref interceptor) = self.auth_interceptor {
            interceptor
                .authorize_request(request, "/charybdis.entities.EntityService/DeleteEntity")?
        } else {
            request
        };

        let entity_id = request.into_inner().id;
        tracing::Span::current().record("entity.id", &entity_id);

        // Fetch entity before deletion for event metadata
        let entity = self
            .repository
            .get_by_id(&entity_id)
            .await
            .map_err(|e| {
                error!(entity_id = %entity_id, error = %e, "Failed to get entity for deletion");
                Status::from(e)
            })?
            .ok_or_else(|| {
                Status::not_found(format!("Entity not found: {}", entity_id))
            })?;

        let entity_kind = entity.kind.clone();
        tracing::Span::current().record("entity.kind", &entity_kind);

        let deleted = self.repository.delete(&entity_id).await.map_err(|e| {
            error!(entity_id = %entity_id, error = %e, "Failed to delete entity");
            Status::from(e)
        })?;

        if !deleted {
            return Err(Status::not_found(format!("Entity not found: {}", entity_id)));
        }

        info!(entity_id = %entity_id, entity_kind = %entity_kind, "Entity deleted");

        if let Ok(uuid) = Uuid::parse_str(&entity_id) {
            let event = EntityEvent::deleted(uuid)
                .with_metadata("entity_kind".to_string(), entity_kind.clone());

            if let Err(e) = self.event_bus.publish(event).await {
                error!(entity_id = %entity_id, error = %e, "Failed to publish entity deleted event");
            }
        }

        let duration = start.elapsed().as_secs_f64();
        self.metrics
            .record_entity_operation("delete", &entity_kind, duration);

        Ok(Response::new(DeleteEntityResponse { success: true }))
    }

    #[instrument(skip(self, request))]
    async fn list_entities(
        &self,
        request: Request<ListEntitiesRequest>,
    ) -> Result<Response<ListEntitiesResponse>, Status> {
        let start = Instant::now();

        let request = if let Some(ref interceptor) = self.auth_interceptor {
            interceptor
                .authorize_request(request, "/charybdis.entities.EntityService/ListEntities")?
        } else {
            request
        };

        let req = request.into_inner();
        let kind = if req.kind.is_empty() { None } else { Some(req.kind.as_str()) };
        let name = if req.name.is_empty() { None } else { Some(req.name.as_str()) };
        let page_size = if req.page_size == 0 { 100 } else { req.page_size };
        let page_token = if req.page_token.is_empty() { None } else { Some(req.page_token.as_str()) };

        let paginated = self
            .repository
            .list_paginated(kind, name, page_size, page_token)
            .await
            .map_err(|e| {
                error!(error = %e, "Failed to list entities");
                Status::from(e)
            })?;

        let duration = start.elapsed().as_secs_f64();
        self.metrics
            .record_entity_operation("list", kind.unwrap_or("all"), duration);

        Ok(Response::new(ListEntitiesResponse {
            entities: paginated.entities,
            next_page_token: paginated.next_page_token.unwrap_or_default(),
            total_count: paginated.total_count,
        }))
    }
}

const VALID_KINDS: &[&str] = &[
    "Service",
    "System",
    "Component",
    "API",
    "User",
    "Group",
    "Domain",
    "Resource",
    "Finding",
];

fn validate_entity(entity: &charybdis::entities::Entity) -> Result<(), Status> {
    use charybdis::entities::entity::Metadata;

    if entity.kind.is_empty() {
        return Err(Status::invalid_argument("Entity kind is required"));
    }
    if !VALID_KINDS.contains(&entity.kind.as_str()) {
        return Err(Status::invalid_argument(format!(
            "Invalid entity kind '{}'. Valid kinds: {}",
            entity.kind,
            VALID_KINDS.join(", ")
        )));
    }

    let metadata = entity
        .metadata
        .as_ref()
        .ok_or_else(|| Status::invalid_argument("Entity metadata is required"))?;

    // Verify metadata variant matches declared kind
    let metadata_kind = match metadata {
        Metadata::ServiceMetadata(_) => "Service",
        Metadata::SystemMetadata(_) => "System",
        Metadata::ComponentMetadata(_) => "Component",
        Metadata::ApiMetadata(_) => "API",
        Metadata::UserMetadata(_) => "User",
        Metadata::GroupMetadata(_) => "Group",
        Metadata::DomainMetadata(_) => "Domain",
        Metadata::ResourceMetadata(_) => "Resource",
        Metadata::FindingMetadata(_) => "Finding",
        _ => return Ok(()), // Plugin metadata — skip kind check
    };

    if metadata_kind != entity.kind {
        return Err(Status::invalid_argument(format!(
            "Metadata type '{}' does not match entity kind '{}'",
            metadata_kind, entity.kind
        )));
    }

    Ok(())
}