From 6e7bfd16f3c2824b93b192e688acb0aa39a2fcde Mon Sep 17 00:00:00 2001 From: Bob McWhirter Date: Fri, 15 Mar 2024 11:59:50 -0400 Subject: [PATCH] Break `description_en` into another table with a `lang` key to be more flexible. Test the aforementioned. --- entity/src/lib.rs | 1 + entity/src/vulnerability.rs | 11 ++- entity/src/vulnerability_description.rs | 29 +++++++ graph/src/graph/vulnerability.rs | 38 +++++++-- ingestors/src/cve/loader.rs | 31 +++++--- migration/src/lib.rs | 2 + .../src/m0000011_create_vulnerability.rs | 2 - ...000012_create_vulnerability_description.rs | 77 +++++++++++++++++++ 8 files changed, 168 insertions(+), 23 deletions(-) create mode 100644 entity/src/vulnerability_description.rs create mode 100644 migration/src/m0000012_create_vulnerability_description.rs diff --git a/entity/src/lib.rs b/entity/src/lib.rs index 67d2b5d3a..a530ae9f9 100644 --- a/entity/src/lib.rs +++ b/entity/src/lib.rs @@ -16,3 +16,4 @@ pub mod sbom; pub mod sbom_describes_cpe22; pub mod sbom_describes_package; pub mod vulnerability; +pub mod vulnerability_description; diff --git a/entity/src/vulnerability.rs b/entity/src/vulnerability.rs index 16d351f19..390c454ca 100644 --- a/entity/src/vulnerability.rs +++ b/entity/src/vulnerability.rs @@ -1,4 +1,4 @@ -use crate::advisory_vulnerability; +use crate::{advisory_vulnerability, vulnerability_description}; use sea_orm::entity::prelude::*; #[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)] @@ -8,13 +8,14 @@ pub struct Model { pub id: i32, pub identifier: String, pub title: Option, - pub description_en: Option, } #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] pub enum Relation { #[sea_orm(has_many = "super::advisory_vulnerability::Entity")] AdvisoryVulnerabilities, + #[sea_orm(has_many = "super::vulnerability_description::Entity")] + Descriptions, } impl Related for Entity { @@ -23,4 +24,10 @@ impl Related for Entity { } } +impl Related for Entity { + fn to() -> RelationDef { + Relation::Descriptions.def() + } +} + impl ActiveModelBehavior for ActiveModel {} diff --git a/entity/src/vulnerability_description.rs b/entity/src/vulnerability_description.rs new file mode 100644 index 000000000..e93d4b4e5 --- /dev/null +++ b/entity/src/vulnerability_description.rs @@ -0,0 +1,29 @@ +use crate::vulnerability; +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)] +#[sea_orm(table_name = "vulnerability_description")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub vulnerability_id: i32, + pub lang: String, + pub description: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "super::vulnerability::Entity", + from = "super::vulnerability_description::Column::VulnerabilityId" + to = "super::vulnerability::Column::Id")] + Vulnerability, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Vulnerability.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/graph/src/graph/vulnerability.rs b/graph/src/graph/vulnerability.rs index d2de94eac..469eca52e 100644 --- a/graph/src/graph/vulnerability.rs +++ b/graph/src/graph/vulnerability.rs @@ -5,13 +5,15 @@ use crate::graph::error::Error; use crate::graph::Graph; use sea_orm::ActiveValue::{Set, Unchanged}; use sea_orm::{ - ActiveModelTrait, ColumnTrait, EntityTrait, NotSet, QueryFilter, QuerySelect, RelationTrait, + ActiveModelTrait, ColumnTrait, EntityTrait, ModelTrait, NotSet, QueryFilter, QuerySelect, + RelationTrait, }; use sea_query::JoinType; use std::fmt::{Debug, Formatter}; use trustify_common::db::Transactional; +use trustify_entity as entity; use trustify_entity::vulnerability::Model; -use trustify_entity::{advisory, advisory_vulnerability, vulnerability}; +use trustify_entity::{advisory, advisory_vulnerability, vulnerability, vulnerability_description}; impl Graph { pub async fn ingest_vulnerability( @@ -26,7 +28,6 @@ impl Graph { id: Default::default(), identifier: Set(identifier.to_string()), title: NotSet, - description_en: NotSet, }; Ok((self, entity.insert(&self.connection(tx)).await?).into()) @@ -95,18 +96,39 @@ impl VulnerabilityContext { Ok(()) } - pub async fn set_description_en( + pub async fn add_description( &self, - description_en: Option, + lang: String, + description: String, tx: Transactional<'_>, ) -> Result<(), Error> { - let mut entity: vulnerability::ActiveModel = self.vulnerability.clone().into(); - entity.description_en = Set(description_en); + let model = vulnerability_description::ActiveModel { + id: Default::default(), + vulnerability_id: Set(self.vulnerability.id), + lang: Set(lang), + description: Set(description), + }; - entity.save(&self.graph.connection(tx)).await?; + model.save(&self.graph.connection(tx)).await?; Ok(()) } + + pub async fn descriptions( + &self, + lang: &str, + tx: Transactional<'_>, + ) -> Result, Error> { + Ok(self + .vulnerability + .find_related(entity::vulnerability_description::Entity) + .filter(vulnerability_description::Column::Lang.eq(lang)) + .all(&self.graph.connection(tx)) + .await? + .drain(..) + .map(|e| e.description) + .collect()) + } } #[cfg(test)] diff --git a/ingestors/src/cve/loader.rs b/ingestors/src/cve/loader.rs index 606579bd2..7bebae9be 100644 --- a/ingestors/src/cve/loader.rs +++ b/ingestors/src/cve/loader.rs @@ -41,17 +41,15 @@ impl<'g> CveLoader<'g> { .set_title(cve.containers.cna.title.clone(), Transactional::Some(&tx)) .await?; - vulnerability - .set_description_en( - cve.containers - .cna - .descriptions - .iter() - .find(|e| e.lang == "en") - .map(|en| en.value.clone()), - Transactional::Some(&tx), - ) - .await?; + for description in cve.containers.cna.descriptions { + vulnerability + .add_description( + description.lang, + description.value, + Transactional::Some(&tx), + ) + .await?; + } let hashes = reader.hashes(); let sha256 = hex::encode(hashes.sha256.as_ref()); @@ -128,6 +126,17 @@ mod test { assert!(loaded_advisory.is_some()); + let loaded_vulnerability = loaded_vulnerability.unwrap(); + + let descriptions = loaded_vulnerability + .descriptions("en", Transactional::None) + .await?; + + assert_eq!(1, descriptions.len()); + + assert!(descriptions[0] + .starts_with("Canarytokens helps track activity and actions on a network")); + Ok(()) } } diff --git a/migration/src/lib.rs b/migration/src/lib.rs index 316f80cc1..1b5c54baa 100644 --- a/migration/src/lib.rs +++ b/migration/src/lib.rs @@ -18,6 +18,7 @@ mod m0000190_sbom_describes_package; mod m0000200_create_relationship; mod m0000210_create_package_relates_to_package; +mod m0000012_create_vulnerability_description; mod m0000035_create_cpe22; mod m0000220_create_qualified_package_transitive_function; mod m0000230_create_importer; @@ -30,6 +31,7 @@ impl MigratorTrait for Migrator { vec![ Box::new(m0000010_create_sbom::Migration), Box::new(m0000011_create_vulnerability::Migration), + Box::new(m0000012_create_vulnerability_description::Migration), Box::new(m0000030_create_advisory::Migration), Box::new(m0000032_create_advisory_vulnerability::Migration), Box::new(m0000040_create_package::Migration), diff --git a/migration/src/m0000011_create_vulnerability.rs b/migration/src/m0000011_create_vulnerability.rs index ca1a93fe8..50306ee78 100644 --- a/migration/src/m0000011_create_vulnerability.rs +++ b/migration/src/m0000011_create_vulnerability.rs @@ -32,7 +32,6 @@ impl MigrationTrait for Migration { .not_null(), ) .col(ColumnDef::new(Vulnerability::Title).string()) - .col(ColumnDef::new(Vulnerability::DescriptionEn).string()) .to_owned(), ) .await @@ -53,5 +52,4 @@ pub enum Vulnerability { // -- Identifier, Title, - DescriptionEn, } diff --git a/migration/src/m0000012_create_vulnerability_description.rs b/migration/src/m0000012_create_vulnerability_description.rs new file mode 100644 index 000000000..e0eb2e0f5 --- /dev/null +++ b/migration/src/m0000012_create_vulnerability_description.rs @@ -0,0 +1,77 @@ +use crate::m0000011_create_vulnerability::Vulnerability; +use crate::m0000032_create_advisory_vulnerability::AdvisoryVulnerability; +use sea_orm_migration::prelude::*; + +use crate::Now; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + // Replace the sample below with your own migration scripts + manager + .create_table( + Table::create() + .table(VulnerabilityDescription::Table) + .if_not_exists() + .col( + ColumnDef::new(VulnerabilityDescription::Id) + .integer() + .not_null() + .auto_increment() + .primary_key(), + ) + .col( + ColumnDef::new(AdvisoryVulnerability::VulnerabilityId) + .integer() + .not_null(), + ) + .foreign_key( + ForeignKey::create() + .from_col(VulnerabilityDescription::VulnerabilityId) + .to(Vulnerability::Table, Vulnerability::Id) + .on_delete(ForeignKeyAction::Cascade), + ) + .col( + ColumnDef::new(VulnerabilityDescription::Timestamp) + .timestamp_with_time_zone() + .default(Func::cust(Now)), + ) + .col( + ColumnDef::new(VulnerabilityDescription::Lang) + .string() + .not_null(), + ) + .col( + ColumnDef::new(VulnerabilityDescription::Description) + .string() + .not_null(), + ) + .to_owned(), + ) + .await + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .drop_table( + Table::drop() + .table(VulnerabilityDescription::Table) + .to_owned(), + ) + .await + } +} + +#[derive(DeriveIden)] +pub enum VulnerabilityDescription { + Table, + Id, + Timestamp, + // -- + VulnerabilityId, + Lang, + Description, +}