feat: Implemented model for dependency vulnerabilities, tested

3 files changed, 26 insertions, 52 deletions

diff --git a/sec_certs/dataset/common_criteria.py b/sec_certs/dataset/common_criteria.py
index 09dcbe23..a05b341e 100644
--- a/sec_certs/dataset/common_criteria.py
+++ b/sec_certs/dataset/common_criteria.py
@@ -19,6 +19,7 @@ from sec_certs.config.configuration import config
 from sec_certs.dataset.dataset import Dataset, logger
 from sec_certs.dataset.protection_profile import ProtectionProfileDataset
 from sec_certs.model.dependency_finder import DependencyFinder
+from sec_certs.model.dependency_vulnerability_finder import DependencyVulnerabilityFinder
 from sec_certs.sample.cc_maintenance_update import CommonCriteriaMaintenanceUpdate
 from sec_certs.sample.common_criteria import CommonCriteriaCert
 from sec_certs.sample.protection_profile import ProtectionProfile
@@ -685,8 +686,13 @@ class CCDataset(Dataset[CommonCriteriaCert], ComplexSerializableType):
             cert.compute_heuristics_cert_id(self.all_cert_ids)
 
     def _compute_dependency_vulnerabilities(self):
-        for cert in self:
-            cert.find_certificate_dependency_vulnerabilities(self)
+        cve_dependency_finder = DependencyVulnerabilityFinder(self.certs)
+        cve_dependency_finder.fit()
+
+        for dgst in self.certs:
+            dependency_cve = cve_dependency_finder.get_dependency_cve(dgst)
+            self.certs[dgst].heuristics.direct_dependency_cves = dependency_cve.direct_dependency_cves
+            self.certs[dgst].heuristics.indirect_dependency_cves = dependency_cve.indirect_dependency_cves
 
     def _compute_heuristics(self, use_nist_cpe_matching_dict: bool = True) -> None:
         self._compute_cert_labs()
diff --git a/sec_certs/model/dependency_vulnerability_finder.py b/sec_certs/model/dependency_vulnerability_finder.py
index 4ea9130b..5eac59a8 100644
--- a/sec_certs/model/dependency_vulnerability_finder.py
+++ b/sec_certs/model/dependency_vulnerability_finder.py
@@ -19,7 +19,7 @@ class DependencyCVE(ComplexSerializableType):
 
 
 Certificates = Dict[str, Certificate]
-Vulnerabilities = Dict[str, Optional[Set[str]]]
+Vulnerabilities = Dict[str, Dict[str, Optional[Set[str]]]]
 
 
 class DependencyVulnerabilityFinder:
@@ -57,7 +57,7 @@ class DependencyVulnerabilityFinder:
         for cert_id in references:
             cert_id_occurrences = dataset_cert_id_occurrences.get(cert_id, None)
 
-            if cert_id is None or cert_id_occurrences[cert_id] >= 2:
+            if cert_id is None or cert_id_occurrences >= 2:
                 continue
 
             for dgst in self.certificates:
@@ -75,12 +75,18 @@ class DependencyVulnerabilityFinder:
         for dgst in self.certificates:
             cert_id = self.certificates[dgst].heuristics.cert_id
 
-            if cert_id is None or cert_id_occurrences[cert_id] >= 2:
+            if cert_id is None:
+                continue
+
+            if cert_id_occurrences[cert_id] >= 2:
                 thrown_away_counter += 1
                 continue
 
-            self.vulnerabilities[dgst][DependencyType.DIRECT] = self._get_dependency_cves(dgst, DependencyType.DIRECT)
-            self.vulnerabilities[dgst][DependencyType.INDIRECT] = self._get_dependency_cves(
+            self.vulnerabilities[dgst] = {}
+            self.vulnerabilities[dgst][DependencyType.DIRECT.value] = self._get_dependency_cves(
+                dgst, DependencyType.DIRECT
+            )
+            self.vulnerabilities[dgst][DependencyType.INDIRECT.value] = self._get_dependency_cves(
                 dgst, DependencyType.INDIRECT
             )
 
@@ -89,8 +95,11 @@ class DependencyVulnerabilityFinder:
 
         return self.vulnerabilities
 
-    def get_dependency_vulnerabilities(self, dgst: str) -> DependencyCVE:
+    def get_dependency_cve(self, dgst: str) -> DependencyCVE:
+        if not self.vulnerabilities.get(dgst):
+            return DependencyCVE(direct_dependency_cves=None, indirect_dependency_cves=None)
+
         return DependencyCVE(
-            self.vulnerabilities[dgst][DependencyType.DIRECT],
-            self.vulnerabilities[dgst][DependencyType.INDIRECT],
+            self.vulnerabilities[dgst][DependencyType.DIRECT.value],
+            self.vulnerabilities[dgst][DependencyType.INDIRECT.value],
         )
diff --git a/sec_certs/sample/common_criteria.py b/sec_certs/sample/common_criteria.py
index 00dbca3a..9dd86c28 100644
--- a/sec_certs/sample/common_criteria.py
+++ b/sec_certs/sample/common_criteria.py
@@ -1,12 +1,11 @@
 import copy
-import logging
 import operator
 from dataclasses import dataclass, field
 from datetime import date, datetime
 from enum import Enum
 from functools import partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Set, Tuple, Union
 
 import requests
 from bs4 import Tag
@@ -20,10 +19,6 @@ from sec_certs.sample.protection_profile import ProtectionProfile
 from sec_certs.serialization.json import ComplexSerializableType
 from sec_certs.serialization.pandas import PandasSerializableType
 
-if TYPE_CHECKING:
-    from sec_certs.dataset.common_criteria import CCDataset
-
-
 HEADERS = {
     "anssi": helpers.search_only_headers_anssi,
     "bsi": helpers.search_only_headers_bsi,
@@ -697,42 +692,6 @@ class CommonCriteriaCert(
         self.heuristics.cert_id = self.pdf_data.cert_id
         self.normalize_cert_id(all_cert_ids)
 
-    def _get_dependency_cves(self, dset: "CCDataset", dependency_type: DependencyType) -> Optional[Set[str]]:
-        dependency_type_dict = {
-            DependencyType.DIRECT.value: self.heuristics.report_references.directly_referenced_by,
-            DependencyType.INDIRECT.value: self.heuristics.report_references.indirectly_referenced_by,
-        }
-
-        references = dependency_type_dict[dependency_type.value]
-
-        if not references:
-            return None
-
-        vulnerabilities = set()
-        dataset_cert_id_occurrences = dset.get_dataset_cert_ids_occurrences()
-        thrown_away_counter = 0
-
-        for cert_id in references:
-            cert_id_occurrence = dataset_cert_id_occurrences.get(cert_id, None)
-
-            # Not in dataset or more than 2 occurrences in dataset -> skip
-            if cert_id_occurrence is None or cert_id_occurrence >= 2:
-                thrown_away_counter += 1
-                continue
-
-            for cert_obj in dset:
-                if cert_obj.heuristics.cert_id == cert_id and cert_obj.heuristics.related_cves:
-                    vulnerabilities.update(cert_obj.heuristics.related_cves)
-
-        if thrown_away_counter >= 1:
-            logging.warning("There were total of %s certificates skipped due to duplicity", thrown_away_counter)
-
-        return vulnerabilities if vulnerabilities else None
-
-    def find_certificate_dependency_vulnerabilities(self, dset: "CCDataset"):
-        self.heuristics.direct_dependency_cves = self._get_dependency_cves(dset, DependencyType.DIRECT)
-        self.heuristics.indirect_dependency_cves = self._get_dependency_cves(dset, DependencyType.INDIRECT)
-
     @staticmethod
     def _is_anssi_cert(cert_id: str) -> bool:
         return cert_id.startswith("ANSS")