finalize cpe matching for on/with configurations

4 files changed, 107 insertions, 135 deletions

diff --git a/src/sec_certs/dataset/cve.py b/src/sec_certs/dataset/cve.py
index 65d141c3..43cd5b80 100644
--- a/src/sec_certs/dataset/cve.py
+++ b/src/sec_certs/dataset/cve.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import collections
 import datetime
 import glob
 import itertools
@@ -57,7 +58,7 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType):
 
     def _filter_cves_with_cpe_configurations(self) -> None:
         """
-        Method filters the subset of CVEs, which contain at least one CPE configuration.
+        Method filters the subset of CVE dataset thah contain at least one CPE configuration in the CVE.
         """
         self.cves_with_vulnerable_configurations = [cve for cve in self if cve.vulnerable_cpe_configurations]
 
@@ -83,7 +84,7 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType):
         for cve in tqdm(self, desc="Building-up lookup dictionaries for fast CVE matching"):
             # See note above, we use matching_dict.get(cpe, []) instead of matching_dict[cpe] as would be expected
             if use_nist_mapping:
-                vulnerable_configurations = set(
+                vulnerable_configurations = list(
                     itertools.chain.from_iterable(matching_dict.get(cpe, []) for cpe in cve.vulnerable_cpes)
                 )
             else:
@@ -133,7 +134,6 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType):
             cls.download_cves(tmp_dir, start_year, end_year)
             json_files = glob.glob(tmp_dir + "/*.json")
 
-            all_cves = {}
             logger.info("Downloaded required resources. Building CVEDataset from jsons.")
             results = process_parallel(
                 cls.from_nist_json,
@@ -141,38 +141,30 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType):
                 use_threading=False,
                 progress_bar_desc="Building CVEDataset from jsons",
             )
-            for r in results:
-                all_cves.update(r.cves)
-
-        return cls(all_cves, json_path)
+        return cls(dict(collections.ChainMap(*(x.cves for x in results))), json_path)
 
     def _get_cve_ids_for_cpe_uri(self, cpe_uri: str) -> set[str]:
         return self.cpe_to_cve_ids_lookup.get(cpe_uri, set())
 
-    def _get_cves_from_exactly_matched_cpes(self, cpe_matches: set[str]) -> set[str]:
-        return set(itertools.chain.from_iterable([self._get_cve_ids_for_cpe_uri(cpe_uri) for cpe_uri in cpe_matches]))
-
-    def _get_cves_from_cpe_configurations(self, cpe_matches: set[str]) -> set[str]:
-        def do_cve_configurations_match_cpe_matches(cve: CVE, cpe_matches: set[str]) -> bool:
-            return any(
-                [cpe_configuration.match(cpe_matches) for cpe_configuration in cve.vulnerable_cpe_configurations]
-            )
+    def _get_cves_from_exactly_matched_cpes(self, cpe_uris: set[str]) -> set[str]:
+        return set(itertools.chain.from_iterable([self._get_cve_ids_for_cpe_uri(cpe_uri) for cpe_uri in cpe_uris]))
 
+    def _get_cves_from_cpe_configurations(self, cpe_uris: set[str]) -> set[str]:
         return {
             cve.cve_id
             for cve in self.cves_with_vulnerable_configurations
-            if do_cve_configurations_match_cpe_matches(cve, cpe_matches)
+            if any(configuration.matches(cpe_uris) for configuration in cve.vulnerable_cpe_configurations)
         }
 
-    def get_cves_from_matched_cpes(self, cpe_matches: set[str]) -> set[str]:
+    def get_cves_from_matched_cpes(self, cpe_uris: set[str]) -> set[str]:
         """
         Method returns the set of CVEs which are matched to the set of CPEs.
         First are matched the classic CPEs to CVEs with lookup dict and then are matched the
         'AND' type CPEs containing platform.
         """
         return {
-            *self._get_cves_from_exactly_matched_cpes(cpe_matches),
-            *self._get_cves_from_cpe_configurations(cpe_matches),
+            *self._get_cves_from_exactly_matched_cpes(cpe_uris),
+            *self._get_cves_from_cpe_configurations(cpe_uris),
         }
 
     def filter_related_cpes(self, relevant_cpes: set[CPE]):
@@ -186,7 +178,13 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType):
         cve_ids_to_delete = []
         for cve in self:
             n_cpes_orig = len(cve.vulnerable_cpes)
-            cve.vulnerable_cpes = list(filter(lambda x: x in relevant_cpes, cve.vulnerable_cpes))
+            cve.vulnerable_cpes = [x for x in cve.vulnerable_cpes if x in relevant_cpes]
+            cve.vulnerable_cpe_configurations = [
+                x
+                for x in cve.vulnerable_cpe_configurations
+                if x.platform.uri in relevant_cpes and any(y.uri in relevant_cpes for y in x.cpes)
+            ]
+
             total_deleted_cpes += n_cpes_orig - len(cve.vulnerable_cpes)
             if not cve.vulnerable_cpes:
                 cve_ids_to_delete.append(cve.cve_id)
diff --git a/src/sec_certs/dataset/dataset.py b/src/sec_certs/dataset/dataset.py
index 9481d9f7..614f60cf 100644
--- a/src/sec_certs/dataset/dataset.py
+++ b/src/sec_certs/dataset/dataset.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import itertools
 import json
 import logging
 import re
@@ -527,18 +526,14 @@ class Dataset(Generic[CertSubType, AuxiliaryDatasetsSubType], ComplexSerializabl
             )
             return
 
-        relevant_cpes = set(itertools.chain.from_iterable(x.heuristics.cpe_matches for x in cpe_rich_certs))
-        self.auxiliary_datasets.cve_dset.filter_related_cpes(relevant_cpes)
+        # The following lines don't bring any speed-up. They may potentially save memory if rest of CVEs is cleaned explicitly
+        # relevant_cpes = set(itertools.chain.from_iterable(x.heuristics.cpe_matches for x in cpe_rich_certs))
+        # self.auxiliary_datasets.cve_dset.filter_related_cpes(relevant_cpes)
 
         cert: Certificate
         for cert in tqdm(cpe_rich_certs, desc="Computing related CVES"):
-            if cert.heuristics.cpe_matches:
-                related_cves = self.auxiliary_datasets.cve_dset.get_cves_from_matched_cpes(cert.heuristics.cpe_matches)
-
-                if related_cves:
-                    cert.heuristics.related_cves = related_cves
-            else:
-                cert.heuristics.related_cves = None
+            related_cves = self.auxiliary_datasets.cve_dset.get_cves_from_matched_cpes(cert.heuristics.cpe_matches)
+            cert.heuristics.related_cves = related_cves if related_cves else None
 
         n_vulnerable = len([x for x in cpe_rich_certs if x.heuristics.related_cves])
         n_vulnerabilities = sum([len(x.heuristics.related_cves) for x in cpe_rich_certs if x.heuristics.related_cves])
diff --git a/src/sec_certs/sample/cpe.py b/src/sec_certs/sample/cpe.py
index 664e70e3..1e3cbf12 100644
--- a/src/sec_certs/sample/cpe.py
+++ b/src/sec_certs/sample/cpe.py
@@ -10,13 +10,12 @@ from sec_certs.serialization.pandas import PandasSerializableType
 from sec_certs.utils import helpers
 
 
+@dataclass
 class CPEConfiguration(ComplexSerializableType):
     __slots__ = ["platform", "cpes"]
 
-    def __init__(self, platform: str, cpes: set[str]) -> None:
-        super().__init__()
-        self.platform: str = platform
-        self.cpes: set[str] = cpes
+    platform: CPE
+    cpes: list[CPE]
 
     def __hash__(self) -> int:
         return hash(self.platform) + sum([hash(cpe) for cpe in self.cpes])
@@ -25,17 +24,19 @@ class CPEConfiguration(ComplexSerializableType):
         return self.platform < other.platform
 
     def __eq__(self, other: Any) -> bool:
-        return isinstance(other, self.__class__) and self.platform == other.platform and self.cpes == other.cpes
+        return (
+            isinstance(other, self.__class__) and self.platform == other.platform and set(self.cpes) == set(other.cpes)
+        )
 
-    def match(self, set_of_cpes: set[str]) -> bool:
+    def matches(self, other_cpe_uris: set[str]) -> bool:
         """
         For a given set of CPEs method returns boolean if the CPE configuration is
         matched or not.
         """
-        return self.platform in set_of_cpes and any(list(set_of_cpes))
+        return self.platform.uri in other_cpe_uris and any(x.uri in other_cpe_uris for x in self.cpes)
 
 
-@dataclass(init=False)
+@dataclass
 class CPE(PandasSerializableType, ComplexSerializableType):
     uri: str
     version: str
@@ -113,6 +114,8 @@ class CPE(PandasSerializableType, ComplexSerializableType):
     def pandas_tuple(self) -> tuple:
         return self.uri, self.vendor, self.item_name, self.version, self.title
 
+    # We cannot use frozen=True. It does not work with __slots__ prior to Python 3.10 dataclasses
+    # Hence we manually provide __hash__ and __eq__ despite not guaranteeing immutability
     def __hash__(self) -> int:
         return hash((self.uri, self.start_version, self.end_version))
 
diff --git a/src/sec_certs/sample/cve.py b/src/sec_certs/sample/cve.py
index 9b2c2437..2289b0e1 100644
--- a/src/sec_certs/sample/cve.py
+++ b/src/sec_certs/sample/cve.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import datetime
 import itertools
 from dataclasses import dataclass
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Iterable
 
 from dateutil.parser import isoparse
 
@@ -12,9 +12,9 @@ from sec_certs.serialization.json import ComplexSerializableType
 from sec_certs.serialization.pandas import PandasSerializableType
 
 
-@dataclass(init=False)
+@dataclass
 class CVE(PandasSerializableType, ComplexSerializableType):
-    @dataclass(eq=True)
+    @dataclass
     class Impact(ComplexSerializableType):
         base_score: float
         severity: str
@@ -47,8 +47,8 @@ class CVE(PandasSerializableType, ComplexSerializableType):
             raise ValueError("NIST Dict for CVE Impact badly formatted.")
 
     cve_id: str
-    vulnerable_cpes: set[CPE]
-    vulnerable_cpe_configurations: set[CPEConfiguration]
+    vulnerable_cpes: list[CPE]
+    vulnerable_cpe_configurations: list[CPEConfiguration]
     impact: Impact
     published_date: datetime.datetime | None
     cwe_ids: set[str] | None
@@ -66,30 +66,13 @@ class CVE(PandasSerializableType, ComplexSerializableType):
         "cwe_ids",
     ]
 
-    def __init__(
-        self,
-        cve_id: str,
-        vulnerable_cpes: set[CPE],
-        vulnerable_cpe_configurations: set[CPEConfiguration],
-        impact: Impact,
-        published_date: str,
-        cwe_ids: set[str] | None,
-    ):
-        super().__init__()
-        self.cve_id = cve_id
-        self.vulnerable_cpes = vulnerable_cpes
-        self.vulnerable_cpe_configurations = vulnerable_cpe_configurations
-        self.impact = impact
-        self.published_date = isoparse(published_date)
-        self.cwe_ids = cwe_ids
-
+    # We cannot use frozen=True. It does not work with __slots__ prior to Python 3.10 dataclasses
+    # Hence we manually provide __hash__ and __eq__ despite not guaranteeing immutability
     def __hash__(self) -> int:
         return hash(self.cve_id)
 
     def __eq__(self, other: object) -> bool:
-        if not isinstance(other, CVE):
-            return False
-        return self.cve_id == other.cve_id
+        return isinstance(other, CVE) and self.cve_id == other.cve_id
 
     def __lt__(self, other: object) -> bool:
         if not isinstance(other, CVE):
@@ -124,11 +107,35 @@ class CVE(PandasSerializableType, ComplexSerializableType):
             "cwe_ids": self.cwe_ids,
         }
 
+    @classmethod
+    def from_dict(cls, dct: dict[str, Any]) -> CVE:
+        date_to_take = (
+            isoparse(dct["published_date"]) if isinstance(dct["published_date"], str) else dct["published_date"]
+        )
+        return cls(
+            dct["cve_id"],
+            dct["vulnerable_cpes"],
+            dct["vulnerable_cpe_configurations"],
+            dct["impact"],
+            date_to_take,
+            dct["cwe_ids"],
+        )
+
+    @classmethod
+    def from_nist_dict(cls, dct: dict) -> CVE:
+        cve_id = dct["cve"]["CVE_data_meta"]["ID"]
+        impact = cls.Impact.from_nist_dict(dct)
+        published_date = isoparse(dct["publishedDate"])
+        cwe_ids = cls.parse_cwe_data(dct)
+        cpes, cpe_configurations = CVE.get_cpe_data_from_nodes_list(dct["configurations"]["nodes"])
+
+        return cls(cve_id, cpes, cpe_configurations, impact, published_date, cwe_ids)
+
     @staticmethod
-    def _parse_nist_cpe_dicts(lst: list[dict[str, Any]]) -> list[CPE]:
+    def _parse_nist_cpe_dicts(dictionaries: Iterable[dict[str, Any]]) -> list[CPE]:
         cpes: list[CPE] = []
 
-        for x in lst:
+        for x in dictionaries:
             cpe_uri = x["cpe23Uri"]
             version_start: tuple[str, str] | None
             version_end: tuple[str, str] | None
@@ -157,82 +164,51 @@ class CVE(PandasSerializableType, ComplexSerializableType):
         The <parse_only_vulnerable_cpes> parameter specifies if we want to
         parse only vulnerable CPEs or not.
         """
-        cpe_dicts_to_be_parsed = cpe_list
+        return CVE._parse_nist_cpe_dicts(dct for dct in cpe_list if dct["vulnerable"] or not parse_only_vulnerable_cpes)
 
-        if parse_only_vulnerable_cpes:
-            cpe_dicts_to_be_parsed = [dct for dct in cpe_list if dct["vulnerable"]]
+    @staticmethod
+    def parse_cwe_data(dct: dict) -> set[str] | None:
+        descriptions = dct["cve"]["problemtype"]["problemtype_data"][0]["description"]
+        return {x["value"] for x in descriptions} if descriptions else None
 
-        return CVE._parse_nist_cpe_dicts(cpe_dicts_to_be_parsed)
+    @staticmethod
+    def get_cpe_data_from_nodes_list(lst: list) -> tuple[list[CPE], list[CPEConfiguration]]:
+        or_nodes = [x for x in lst if x["operator"] == "OR"]
+        and_nodes = [x for x in lst if x["operator"] == "AND"]
+        return CVE.get_simple_cpes_from_nodes_list(or_nodes), CVE.get_cpe_configurations_from_node_list(and_nodes)
 
-    @classmethod
-    def from_nist_dict(cls, dct: dict) -> CVE:
+    @staticmethod
+    def get_simple_cpes_from_nodes_list(lst: list) -> list[CPE]:
+        return list(
+            itertools.chain.from_iterable(
+                CVE._parse_nist_dict(node["cpe_match"], parse_only_vulnerable_cpes=True) for node in lst
+            )
+        )
+
+    @staticmethod
+    def get_cpe_configurations_from_node_list(lst: list) -> list[CPEConfiguration]:
         """
-        Will load CVE from dictionary defined at https://nvd.nist.gov/feeds/json/cve/1.1
+        Retrieves only running on/with configurations, not the advanced ones.
+        See more at https://nvd.nist.gov/vuln/vulnerability-detail-pages, section `Configurations`
         """
+        configurations = [CVE.get_cpe_confiugration_from_node(x) for x in lst]
+        return [x for x in configurations if x]
 
-        def get_cpe_configurations_from_and_cpe_dict(children: list[dict]) -> list[CPEConfiguration]:
-            configurations: list[CPEConfiguration] = []
-
-            if not children or len(children) != 2:
-                return configurations
-
-            cpes = CVE._parse_nist_dict(children[0]["cpe_match"], True)
-            vulnerable_cpe_uris = {cpe.uri for cpe in cpes}
-
-            if not cpes:
-                return configurations
-
-            # Platform does not have to be vulnerable necessarily
-            platforms = CVE._parse_nist_dict(children[1]["cpe_match"], False)
-
-            return [CPEConfiguration(platform.uri, vulnerable_cpe_uris) for platform in platforms]
-
-        def get_vulnerable_cpes_from_nist_dict(dct: dict) -> list[list]:
-            def get_vulnerable_cpes_and_cpe_configurations(
-                node: dict, cpes: list[CPE], cpe_configurations: list[CPEConfiguration]
-            ) -> tuple[list[CPE], list[CPEConfiguration]]:
-                """
-                Method traverses node of CPE tree and returns the list of CPEs and CPE configuratios,
-                which depends on if the parent node is OR/AND type.
-                """
-                if node["operator"] == "AND":
-                    cpe_configurations.extend(get_cpe_configurations_from_and_cpe_dict(node["children"]))
-                    return cpes, cpe_configurations
-
-                if "children" in node:
-                    for child in node["children"]:
-                        get_vulnerable_cpes_and_cpe_configurations(child, cpes, cpe_configurations)
-
-                if "cpe_match" not in node:
-                    return cpes, cpe_configurations
-
-                candidates = node["cpe_match"]
-                cpes.extend(CVE._parse_nist_dict(candidates, True))
-
-                return cpes, cpe_configurations
-
-            cpes_and_cpe_configurations = [
-                get_vulnerable_cpes_and_cpe_configurations(x, [], []) for x in dct["configurations"]["nodes"]
-            ]
+    @staticmethod
+    def get_cpe_confiugration_from_node(node: dict) -> CPEConfiguration | None:
+        if node["children"]:
+            if len(node["children"]) != 2:
+                return None
 
-            return [list(t) for t in zip(*cpes_and_cpe_configurations)]
+            # Deep variant should have two children, get CPEs from the first one and declare that product, second is platform
+            cpes = CVE._parse_nist_dict(node["children"][0]["cpe_match"], parse_only_vulnerable_cpes=True)
+            platform = CVE._parse_nist_dict(node["children"][1]["cpe_match"], parse_only_vulnerable_cpes=False)
+            return CPEConfiguration(platform[0], cpes)
+        else:
+            # Shallow variant should have exactly 2 matching CPEs, we declare one a platform, second one the vuln. thing
+            cpes = CVE._parse_nist_dict(node["cpe_match"], parse_only_vulnerable_cpes=True)
 
-        cve_id = dct["cve"]["CVE_data_meta"]["ID"]
-        impact = cls.Impact.from_nist_dict(dct)
-        cpe_and_cpe_configurations = get_vulnerable_cpes_from_nist_dict(dct)
-        # There exist CVEs such as (CVE-2022-0177) which are rejected and do not contain any assinged CPEs
-        vulnerable_cpes = (
-            set(itertools.chain.from_iterable(cpe_and_cpe_configurations[0])) if cpe_and_cpe_configurations else set()
-        )
-        vulnerable_cpe_configurations = (
-            set(itertools.chain.from_iterable(cpe_and_cpe_configurations[1])) if cpe_and_cpe_configurations else set()
-        )
-        published_date = dct["publishedDate"]
-        cwe_ids = cls.parse_cwe_data(dct)
+            if len(cpes) != 2:
+                return None
 
-        return cls(cve_id, vulnerable_cpes, vulnerable_cpe_configurations, impact, published_date, cwe_ids)
-
-    @staticmethod
-    def parse_cwe_data(dct: dict) -> set[str] | None:
-        descriptions = dct["cve"]["problemtype"]["problemtype_data"][0]["description"]
-        return {x["value"] for x in descriptions} if descriptions else None
+            return CPEConfiguration(cpes[0], [cpes[1]])