AIBOM labels

Signed-off-by: desmax74 <mdessi@redhat.com>

Author: desmax74

modules/ingestor/src/service/sbom/cyclonedx.rs

@@ -1,9 +1,11 @@
+
 use crate::{
     graph::{Graph, Outcome, sbom::cyclonedx},
     model::IngestResult,
     service::{Error, Warnings},
 };
 use sea_orm::TransactionTrait;
+use serde_cyclonedx::cyclonedx::v_1_6::Component;
 use tracing::instrument;
 use trustify_common::{hashing::Digests, id::Id};
 use trustify_entity::labels::Labels;
@@ -17,6 +19,11 @@ impl<'g> CyclonedxLoader<'g> {
         Self { graph }
     }
 
+    /*pub fn aiType() -> bool{
+
+        false
+   }*/
+
     #[instrument(skip(self, buffer), err(level=tracing::Level::INFO))]
     pub async fn load(
         &self,
@@ -29,7 +36,7 @@ impl<'g> CyclonedxLoader<'g> {
         let cdx: Box<serde_cyclonedx::cyclonedx::v_1_6::CycloneDx> = serde_json::from_slice(buffer)
             .map_err(|err| Error::UnsupportedFormat(format!("Failed to parse: {err}")))?;
 
-        let labels = labels.add("type", "cyclonedx");
+        let labels = labels_with_ai_type_check(cdx.components.clone());
 
         log::info!(
             "Storing - version: {:?}, serialNumber: {:?}",
@@ -75,6 +82,23 @@ impl<'g> CyclonedxLoader<'g> {
     }
 }
 
+
+fn labels_with_ai_type_check(components: Option<Vec<Component>>) -> Labels{
+    match components {
+        Some(vec) => {
+            for component in vec {
+                if component.type_ == "machine-learning-model" {
+                    return  Labels::new().add("type", "cyclonedx").add("ai", "machine-learning-model");
+                }
+            }
+        }
+        None => {
+            return  Labels::new().add("type", "cyclonedx");
+        }
+    }
+    Labels::new().add("type", "cyclonedx")
+}
+
 #[cfg(test)]
 mod test {
     use crate::service::{Cache, IngestorService};
@@ -105,4 +129,28 @@ mod test {
 
         Ok(())
     }
+
+
+    #[test_context(TrustifyContext)]
+    #[test(tokio::test)]
+    async fn ingest_ai_cyclonedx(ctx: &TrustifyContext) -> Result<(), anyhow::Error> {
+        let db = &ctx.db;
+        let graph = Graph::new(db.clone());
+        let data = document_bytes("cyclonedx/ai/nvidia_canary-1b-v2_aibom.json").await?;
+
+        let ingestor = IngestorService::new(graph, ctx.storage.clone(), Default::default());
+
+        ingestor
+            .ingest(
+                &data,
+                Format::CycloneDX,
+                [("type", "cyclonedx"),("ai", "machine-learning-model")],
+                None,
+                Cache::Skip,
+            )
+            .await
+            .expect("must ingest");
+
+        Ok(())
+    }
 }