feat: Partly implemented module, not integrated with dataset, nor tested

1 files changed, 96 insertions, 0 deletions

diff --git a/sec_certs/model/dependency_vulnerability_finder.py b/sec_certs/model/dependency_vulnerability_finder.py
new file mode 100644
index 00000000..4ea9130b
--- /dev/null
+++ b/sec_certs/model/dependency_vulnerability_finder.py
@@ -0,0 +1,96 @@
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, Optional, Set
+
+from sec_certs.sample.certificate import Certificate
+from sec_certs.serialization.json import ComplexSerializableType
+
+
+class DependencyType(Enum):
+    DIRECT = "direct"
+    INDIRECT = "indirect"
+
+
+@dataclass
+class DependencyCVE(ComplexSerializableType):
+    direct_dependency_cves: Optional[Set[str]] = field(default=None)
+    indirect_dependency_cves: Optional[Set[str]] = field(default=None)
+
+
+Certificates = Dict[str, Certificate]
+Vulnerabilities = Dict[str, Optional[Set[str]]]
+
+
+class DependencyVulnerabilityFinder:
+    def __init__(self, certificates: Certificates):
+        self.vulnerabilities: Vulnerabilities = {}
+        self.certificates = certificates
+
+    def _get_dataset_cert_ids(self) -> Dict[str, int]:
+        cert_id_occurrences: Dict[str, int] = {}
+
+        for dgst in self.certificates:
+            cert_id = self.certificates[dgst].heuristics.cert_id
+
+            if cert_id is None:
+                continue
+
+            cert_id_occurrences[cert_id] = cert_id_occurrences.get(cert_id, 0) + 1
+
+        return cert_id_occurrences
+
+    def _get_dependency_cves(self, dgst: str, dependency_type: DependencyType) -> Optional[Set[str]]:
+        dependency_type_dict = {
+            DependencyType.DIRECT: self.certificates[dgst].heuristics.report_references.directly_referenced_by,
+            DependencyType.INDIRECT: self.certificates[dgst].heuristics.report_references.indirectly_referenced_by,
+        }
+
+        references = dependency_type_dict[dependency_type]
+
+        if not references:
+            return None
+
+        vulnerabilities = set()
+        dataset_cert_id_occurrences = self._get_dataset_cert_ids()
+
+        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:
+                continue
+
+            for dgst in self.certificates:
+                cert_obj = self.certificates[dgst]
+
+                if cert_obj.heuristics.cert_id == cert_id and cert_obj.heuristics.related_cves:
+                    vulnerabilities.update(cert_obj.heuristics.related_cves)
+
+        return vulnerabilities if vulnerabilities else None
+
+    def fit(self) -> Vulnerabilities:
+        cert_id_occurrences = self._get_dataset_cert_ids()
+        thrown_away_counter = 0
+
+        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:
+                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(
+                dgst, DependencyType.INDIRECT
+            )
+
+        if thrown_away_counter > 0:
+            logging.warning("There were total of %s certificates skipped due to duplicity", thrown_away_counter)
+
+        return self.vulnerabilities
+
+    def get_dependency_vulnerabilities(self, dgst: str) -> DependencyCVE:
+        return DependencyCVE(
+            self.vulnerabilities[dgst][DependencyType.DIRECT],
+            self.vulnerabilities[dgst][DependencyType.INDIRECT],
+        )