Make CCSchemeDataset an actual dataset.

6 files changed, 174 insertions, 16 deletions

diff --git a/src/sec_certs/dataset/cc.py b/src/sec_certs/dataset/cc.py
index 2720638c..eb0156c2 100644
--- a/src/sec_certs/dataset/cc.py
+++ b/src/sec_certs/dataset/cc.py
@@ -17,10 +17,12 @@ from bs4 import BeautifulSoup, Tag
 import sec_certs.utils.sanitization
 from sec_certs import constants
 from sec_certs.configuration import config
+from sec_certs.dataset.cc_scheme import CCSchemeDataset, EntryType
 from sec_certs.dataset.cpe import CPEDataset
 from sec_certs.dataset.cve import CVEDataset
 from sec_certs.dataset.dataset import AuxiliaryDatasets, Dataset, logger
 from sec_certs.dataset.protection_profile import ProtectionProfileDataset
+from sec_certs.model.cc_matching import CCSchemeMatcher
 from sec_certs.model.reference_finder import ReferenceFinder
 from sec_certs.model.sar_transformer import SARTransformer
 from sec_certs.model.transitive_vulnerability_finder import TransitiveVulnerabilityFinder
@@ -39,6 +41,7 @@ class CCAuxiliaryDatasets(AuxiliaryDatasets):
     cve_dset: CVEDataset | None = None
     pp_dset: ProtectionProfileDataset | None = None
     mu_dset: CCDatasetMaintenanceUpdates | None = None
+    scheme_dset: CCSchemeDataset | None = None
 
 
 class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializableType):
@@ -139,7 +142,7 @@ class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializable
     @property
     def pp_dataset_path(self) -> Path:
         """
-        Returns directory that holds files associated with Protection profiles
+        Returns a path to the dataset of Protection Profiles
         """
         return self.auxiliary_datasets_dir / "pp_dataset.json"
 
@@ -153,10 +156,17 @@ class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializable
     @property
     def mu_dataset_path(self) -> Path:
         """
-        Returns json that holds the datase of maintenance updates
+        Returns a path to the dataset of maintenance updates
         """
         return self.mu_dataset_dir / "maintenance_updates.json"
 
+    @property
+    def scheme_dataset_path(self) -> Path:
+        """
+        Returns a path to the scheme dataset
+        """
+        return self.auxiliary_datasets_dir / "scheme_dataset.json"
+
     BASE_URL: ClassVar[str] = "https://www.commoncriteriaportal.org"
 
     HTML_PRODUCTS_URL = {
@@ -733,11 +743,12 @@ class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializable
     def process_auxiliary_datasets(self, download_fresh: bool = False) -> None:
         """
         Processes all auxiliary datasets needed during computation. On top of base-class processing,
-        CC handles protection profiles and maintenance updates.
+        CC handles protection profiles, maintenance updates and schemes.
         """
         super().process_auxiliary_datasets(download_fresh)
         self.auxiliary_datasets.pp_dset = self.process_protection_profiles(to_download=download_fresh)
         self.auxiliary_datasets.mu_dset = self.process_maintenance_updates(to_download=download_fresh)
+        self.auxiliary_datasets.scheme_dset = self.process_schemes(to_download=download_fresh)
 
     def process_protection_profiles(
         self, to_download: bool = True, keep_metadata: bool = True
@@ -798,6 +809,30 @@ class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializable
 
         return update_dset
 
+    def process_schemes(self, to_download: bool = True, only_schemes: set[str] | None = None) -> CCSchemeDataset:
+        """
+        Downloads or loads from json a dataset of CC scheme data.
+        """
+        logger.info("Processing CC schemes")
+
+        self.auxiliary_datasets_dir.mkdir(parents=True, exist_ok=True)
+
+        if to_download or not self.scheme_dataset_path.exists():
+            scheme_dset = CCSchemeDataset.from_web(only_schemes)
+        else:
+            scheme_dset = CCSchemeDataset.from_json(self.scheme_dataset_path)
+
+        for scheme in scheme_dset:
+            code = scheme["scheme"]
+            certified = scheme.get(EntryType.Certified)
+            if certified:
+                matches = CCSchemeMatcher.match_all(certified, code, self)
+                for dgst, match in matches.items():
+                    self[dgst].heuristics.scheme_data = match
+            # TODO: Archived??
+
+        return scheme_dset
+
 
 class CCDatasetMaintenanceUpdates(CCDataset, ComplexSerializableType):
     """
diff --git a/src/sec_certs/dataset/cc_scheme.py b/src/sec_certs/dataset/cc_scheme.py
index e518665d..4fb6356a 100644
--- a/src/sec_certs/dataset/cc_scheme.py
+++ b/src/sec_certs/dataset/cc_scheme.py
@@ -3,10 +3,13 @@
 from __future__ import annotations
 
 import hashlib
+import logging
 import tempfile
 import warnings
+from datetime import datetime
+from enum import Enum, auto
 from pathlib import Path
-from typing import Any
+from typing import Any, Callable, Mapping
 from urllib.parse import urljoin
 
 import requests
@@ -16,19 +19,134 @@ from requests import Response
 from urllib3.connectionpool import InsecureRequestWarning
 
 from sec_certs import constants
+from sec_certs.dataset.json_path_dataset import JSONPathDataset
+from sec_certs.serialization.json import ComplexSerializableType
 from sec_certs.utils.sanitization import sanitize_navigable_string as sns
 from sec_certs.utils.tqdm import tqdm
 
+logger = logging.getLogger()
 
-class CCSchemeDataset:
+
+class EntryType(Enum):
+    Certified = auto()
+    InEvaluation = auto()
+    Archived = auto()
+
+
+class CCSchemeDataset(JSONPathDataset, ComplexSerializableType):
     """
-    Not really a dataset of data from CC scheme websites.
+    A dataset of data from CC scheme websites.
 
     Each `.get_*` method returns a list of dict entries from the given scheme.
     The entries do not share many keys, but each one has at least some form
     of a product name and most have a vendor/developer/manufacturer field.
     """
 
+    def __init__(self, schemes, json_path: str | Path = constants.DUMMY_NONEXISTING_PATH):
+        self.schemes = schemes
+        self.json_path = Path(json_path)
+
+    @property
+    def serialized_attributes(self) -> list[str]:
+        return ["schemes"]
+
+    def __iter__(self):
+        yield from self.schemes.values()
+
+    def __getitem__(self, scheme: str):
+        return self.schemes.__getitem__(scheme.upper())
+
+    def __setitem__(self, key: str, value):
+        self.schemes.__setitem__(key.upper(), value)
+
+    def __len__(self) -> int:
+        return len(self.schemes)
+
+    def to_dict(self):
+        return {"schemes": self.schemes}
+
+    @classmethod
+    def from_dict(cls, dct: Mapping) -> CCSchemeDataset:
+        return cls(dct["schemes"])
+
+    @classmethod
+    def from_web(cls, only_schemes: set[str] | None = None) -> CCSchemeDataset:
+        methods: dict[str, list[dict[str, EntryType | Callable]]] = {
+            "AU": [{"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_australia_in_evaluation}],
+            "CA": [
+                {"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_canada_in_evaluation},
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_canada_certified},
+            ],
+            "FR": [{"type": EntryType.Certified, "method": CCSchemeDataset.get_france_certified}],
+            "DE": [{"type": EntryType.Certified, "method": CCSchemeDataset.get_germany_certified}],
+            "IN": [
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_india_certified},
+                {"type": EntryType.Archived, "method": CCSchemeDataset.get_india_archived},
+            ],
+            "IT": [
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_italy_certified},
+                {"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_italy_in_evaluation},
+            ],
+            "JP": [
+                {"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_japan_in_evaluation},
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_japan_certified},
+                {"type": EntryType.Archived, "method": CCSchemeDataset.get_japan_archived},
+            ],
+            "MY": [
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_malaysia_certified},
+                {"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_malaysia_in_evaluation},
+            ],
+            "NL": [
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_netherlands_certified},
+                {"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_netherlands_in_evaluation},
+            ],
+            "NO": [
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_norway_certified},
+                {"type": EntryType.Archived, "method": CCSchemeDataset.get_norway_archived},
+            ],
+            "KO": [
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_korea_certified},
+                {"type": EntryType.Archived, "method": CCSchemeDataset.get_korea_archived},
+            ],
+            "SG": [
+                {"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_singapore_in_evaluation},
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_singapore_certified},
+                {"type": EntryType.Archived, "method": CCSchemeDataset.get_singapore_archived},
+            ],
+            "ES": [{"type": EntryType.Certified, "method": CCSchemeDataset.get_spain_certified}],
+            "SE": [
+                {"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_sweden_in_evaluation},
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_sweden_certified},
+                {"type": EntryType.Archived, "method": CCSchemeDataset.get_sweden_archived},
+            ],
+            "TR": [{"type": EntryType.Certified, "method": CCSchemeDataset.get_turkey_certified}],
+            "US": [
+                {"type": EntryType.InEvaluation, "method": CCSchemeDataset.get_usa_in_evaluation},
+                {"type": EntryType.Certified, "method": CCSchemeDataset.get_usa_certified},
+                {"type": EntryType.Archived, "method": CCSchemeDataset.get_usa_archived},
+            ],
+        }
+        schemes = {}
+        for scheme, sources in methods.items():
+            if only_schemes is not None and scheme not in only_schemes:
+                continue
+            timestamp = datetime.now()
+            entry: dict[str | EntryType, Any] = {
+                "scheme": scheme,
+                "timestamp": timestamp,
+            }
+            for source in sources:
+                source_type: EntryType = source["type"]  # type: ignore
+                source_method: Callable = source["method"]  # type: ignore
+                try:
+                    res = source_method()
+                except Exception as e:
+                    logger.warning(f"Error during {scheme} download of {source_type}: {e}")
+                    continue
+                entry[source_type] = res
+            schemes[scheme] = entry
+        return cls(schemes)
+
     @staticmethod
     def _get(url: str, session, **kwargs) -> Response:
         with warnings.catch_warnings():
diff --git a/src/sec_certs/model/cc_matching.py b/src/sec_certs/model/cc_matching.py
index fdb16d70..3cc38a40 100644
--- a/src/sec_certs/model/cc_matching.py
+++ b/src/sec_certs/model/cc_matching.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Iterable, Mapping
+from typing import Any, Iterable, Mapping, Sequence
 
 from sec_certs.configuration import config
 from sec_certs.model.matching import AbstractMatcher
@@ -79,7 +79,7 @@ class CCSchemeMatcher(AbstractMatcher[CCCertificate]):
     @classmethod
     def match_all(
         cls, entries: list[dict[str, Any]], scheme: str, certificates: Iterable[CCCertificate]
-    ) -> dict[str, Mapping]:
+    ) -> dict[str, dict[str, Any]]:
         """
         Match all entries of a given CC scheme to certificates from the dataset.
 
@@ -89,5 +89,5 @@ class CCSchemeMatcher(AbstractMatcher[CCCertificate]):
         :return: A mapping of certificate digests to entries, without duplicates, not all entries may be present.
         """
         certs: list[CCCertificate] = list(filter(lambda cert: cert.scheme == scheme, certificates))
-        matchers = [CCSchemeMatcher(entry, scheme) for entry in entries]
-        return cls._match_certs(matchers, certs, config.cc_matching_threshold)
+        matchers: Sequence[CCSchemeMatcher] = [CCSchemeMatcher(entry, scheme) for entry in entries]
+        return cls._match_certs(matchers, certs, config.cc_matching_threshold)  # type: ignore
diff --git a/src/sec_certs/model/fips_matching.py b/src/sec_certs/model/fips_matching.py
index 8cf3cb9c..53d75c61 100644
--- a/src/sec_certs/model/fips_matching.py
+++ b/src/sec_certs/model/fips_matching.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Iterable, Mapping
+from typing import Iterable, Mapping, Sequence
 
 from sec_certs.configuration import config
 from sec_certs.model.matching import AbstractMatcher
@@ -60,5 +60,6 @@ class FIPSProcessMatcher(AbstractMatcher[FIPSCertificate]):
         :return: A mapping of certificate digests to entries, without duplicates, not all entries may be present.
         """
         certs: list[FIPSCertificate] = list(certificates)
-        matchers = [FIPSProcessMatcher(entry) for entry in snapshot]
-        return cls._match_certs(matchers, certs, config.fips_matching_threshold)
+        matchers: Sequence[FIPSProcessMatcher] = [FIPSProcessMatcher(entry) for entry in snapshot]
+        # mypy is ridiculous
+        return cls._match_certs(matchers, certs, config.fips_matching_threshold)  # type: ignore
diff --git a/src/sec_certs/model/matching.py b/src/sec_certs/model/matching.py
index 6d49f34f..fc71b2ad 100644
--- a/src/sec_certs/model/matching.py
+++ b/src/sec_certs/model/matching.py
@@ -1,7 +1,7 @@
 import typing
 from abc import ABC, abstractmethod
 from heapq import heappop, heappush
-from typing import Generic
+from typing import Any, Generic, Sequence
 
 from rapidfuzz import fuzz
 
@@ -11,6 +11,8 @@ CertSubType = typing.TypeVar("CertSubType", bound=Certificate)
 
 
 class AbstractMatcher(Generic[CertSubType], ABC):
+    entry: Any
+
     @abstractmethod
     def match(self, cert: CertSubType) -> float:
         raise NotImplementedError
@@ -25,7 +27,7 @@ class AbstractMatcher(Generic[CertSubType], ABC):
         )
 
     @staticmethod
-    def _match_certs(matchers, certs, threshold):
+    def _match_certs(matchers: Sequence["AbstractMatcher"], certs: list[CertSubType], threshold: float):
         scores: list[tuple[float, int, int]] = []
         matched_is: set[int] = set()
         matched_js: set[int] = set()
diff --git a/src/sec_certs/sample/cc.py b/src/sec_certs/sample/cc.py
index 9f73aefa..6af192f6 100644
--- a/src/sec_certs/sample/cc.py
+++ b/src/sec_certs/sample/cc.py
@@ -20,9 +20,10 @@ import sec_certs.utils.sanitization
 from sec_certs import constants
 from sec_certs.cert_rules import SARS_IMPLIED_FROM_EAL, cc_rules, rules, security_level_csv_scan
 from sec_certs.sample.cc_certificate_id import canonicalize
-from sec_certs.sample.certificate import Certificate, References, logger
+from sec_certs.sample.certificate import Certificate
 from sec_certs.sample.certificate import Heuristics as BaseHeuristics
 from sec_certs.sample.certificate import PdfData as BasePdfData
+from sec_certs.sample.certificate import References, logger
 from sec_certs.sample.protection_profile import ProtectionProfile
 from sec_certs.sample.sar import SAR
 from sec_certs.serialization.json import ComplexSerializableType
@@ -437,6 +438,7 @@ class CCCertificate(
         extracted_sars: set[SAR] | None = field(default=None)
         direct_transitive_cves: set[str] | None = field(default=None)
         indirect_transitive_cves: set[str] | None = field(default=None)
+        scheme_data: dict[str, Any] | None = field(default=None)
 
         @property
         def serialized_attributes(self) -> list[str]: