Make DependencyFinder handle duplicates and unknowns.

1 files changed, 102 insertions, 36 deletions

diff --git a/sec_certs/model/dependency_finder.py b/sec_certs/model/dependency_finder.py
index d41cee34..c28db350 100644
--- a/sec_certs/model/dependency_finder.py
+++ b/sec_certs/model/dependency_finder.py
@@ -6,10 +6,13 @@ Certificates = Dict[str, Certificate]
 ReferencedByDirect = Dict[str, Set[str]]
 ReferencedByIndirect = Dict[str, Set[str]]
 Dependencies = Dict[str, Dict[str, Optional[Set[str]]]]
+IDMapping = Dict[str, List[str]]
+UnknownReferences = Dict[str, Set[str]]
 IDLookupFunc = Callable[[Certificate], str]
 ReferenceLookupFunc = Callable[[Certificate], Set[str]]
 
 
+# TODO: All of this can and should be rewritten on top of networkx or some other graph library.
 class DependencyFinder:
     """
     The class assigns references of other certificate instances for each instance.
@@ -19,26 +22,48 @@ class DependencyFinder:
 
     def __init__(self):
         self.dependencies: Dependencies = {}
+        self.id_mapping: IDMapping = {}
         self._fitted: bool = False
 
-    def _add_direct_reference(self, referenced_by: ReferencedByDirect, cert_id: str, this_cert_id: str) -> None:
-        if cert_id not in referenced_by:
-            referenced_by[cert_id] = set()
-        if this_cert_id not in referenced_by[cert_id]:
-            referenced_by[cert_id].add(this_cert_id)
+    def _create_id_mapping(self, certificates: Certificates, id_func: IDLookupFunc) -> None:
+        """
+        Create the ID mapping of certificate IDs to certificate digests.
+
+        Necessary for handling duplicates.
+        """
+        # Create a mapping of certificate ID to certificate digests with that ID.
+        for dgst in certificates:
+            cert_id = id_func(certificates[dgst])
+            c_list = self.id_mapping.setdefault(cert_id, [])
+            c_list.append(dgst)
 
-    def _process_references(
-        self, referenced_by: ReferencedByDirect, referenced_by_indirect: ReferencedByIndirect
-    ) -> None:
+        # Sort digests in ID mapping to have deterministic behavior.
+        # The certificate with the first digest will be used with that ID, others will be discarded.
+        for digests in self.id_mapping.values():
+            digests.sort()
+
+    def _compute_indirect_references(self, referenced_by: ReferencedByDirect) -> ReferencedByIndirect:
+        """
+        Compute indirect references via a BFS algorithm.
+        """
+        referenced_by_indirect: ReferencedByIndirect = {}
+
+        # Populate with direct references.
+        certs_id_list = referenced_by.keys()
+        for cert_id in certs_id_list:
+            referenced_by_indirect[cert_id] = set()
+            for item in referenced_by[cert_id]:
+                referenced_by_indirect[cert_id].add(item)
+
+        # Flood in the indirect ones.
         new_change_detected = True
         while new_change_detected:
             new_change_detected = False
-            certs_id_list = referenced_by.keys()
 
             for cert_id in certs_id_list:
                 tmp_referenced_by_indirect_nums = referenced_by_indirect[cert_id].copy()
                 for referencing in tmp_referenced_by_indirect_nums:
-                    if referencing in referenced_by.keys():
+                    if referencing in certs_id_list:
                         tmp_referencing = referenced_by_indirect[referencing].copy()
                         newly_discovered_references = [
                             x for x in tmp_referencing if x not in referenced_by_indirect[cert_id]
@@ -46,34 +71,33 @@ class DependencyFinder:
                         referenced_by_indirect[cert_id].update(newly_discovered_references)
                         if newly_discovered_references:
                             new_change_detected = True
+        return referenced_by_indirect
 
     def _build_referenced_by(
-        self, certificates: Certificates, id_func: IDLookupFunc, ref_lookup_func: ReferenceLookupFunc
+        self, certificates: Certificates, ref_lookup_func: ReferenceLookupFunc
     ) -> Tuple[ReferencedByDirect, ReferencedByIndirect]:
         referenced_by: ReferencedByDirect = {}
 
-        for cert_obj in certificates.values():
+        for this_cert_id, cert_digests in self.id_mapping.items():
+            # Take the first certificate digest from the ID mapping (to ensure deterministic behavior and resolve duplicates).
+            # TODO: A better approach for handling duplicates in the future would be nice.
+            cert_dgst = cert_digests[0]
+            cert_obj = certificates[cert_dgst]
+
             refs = ref_lookup_func(cert_obj)
             if refs is None:
                 continue
 
-            this_cert_id = id_func(cert_obj)
-            if this_cert_id is None:
-                continue
-
-            # Direct reference
+            # Process direct reference
+            # All are added here, the unknown ones are filtered later on.
             for cert_id in refs:
-                if cert_id != this_cert_id:
-                    self._add_direct_reference(referenced_by, cert_id, this_cert_id)
+                if cert_id == this_cert_id:
+                    continue
+                referenced_by.setdefault(cert_id, set())
+                referenced_by[cert_id].add(this_cert_id)
 
-        referenced_by_indirect: ReferencedByIndirect = {}
-
-        for cert_id in referenced_by.keys():
-            referenced_by_indirect[cert_id] = set()
-            for item in referenced_by[cert_id]:
-                referenced_by_indirect[cert_id].add(item)
-
-        self._process_references(referenced_by, referenced_by_indirect)
+        # Now do the indirect ones
+        referenced_by_indirect = self._compute_indirect_references(referenced_by)
         return referenced_by, referenced_by_indirect
 
     def _get_reverse_dependencies(
@@ -123,23 +147,64 @@ class DependencyFinder:
         """
         if self._fitted:
             raise ValueError("Finder already fitted")
-        referenced_by_direct, referenced_by_indirect = self._build_referenced_by(certificates, id_func, ref_lookup_func)
+        # Create the ID mapping first so that we can resolve duplicates.
+        self._create_id_mapping(certificates, id_func)
+
+        # Build the referenced_by first
+        referenced_by_direct, referenced_by_indirect = self._build_referenced_by(certificates, ref_lookup_func)
 
+        # Build the referencing second (this actually writes into self.dependencies).
         self._build_referencing(certificates, id_func, referenced_by_direct, referenced_by_indirect)
         self._fitted = True
 
-    def predict_single_cert(self, dgst: str) -> References:
+    @property
+    def unknown_references(self) -> UnknownReferences:
+        """
+        Get the unknown references in the fitted dataset (to unknown certificate IDs, not in the dataset during fit).
+        """
+        if not self._fitted:
+            return {}
+        result = {}
+        for cert_id, digests in self.id_mapping.items():
+            cert_digest = digests[0]
+            cert_references = self.dependencies[cert_digest]
+            direct_refs = cert_references["directly_referencing"]
+            if not direct_refs:
+                continue
+            unknowns = set(filter(lambda refd_cert_id: refd_cert_id not in self.id_mapping, direct_refs))
+            if unknowns:
+                result[cert_id] = unknowns
+        return result
+
+    @property
+    def duplicates(self) -> IDMapping:
+        """
+        Get the duplicates in the fitted dataset.
+
+        :return IDMapping: Mapping of certificate ID to digests that share it.
+        """
+        if not self._fitted:
+            return {}
+        return {cert_id: digests for cert_id, digests in self.id_mapping.items() if len(digests) > 1}
+
+    def predict_single_cert(self, dgst: str, keep_unknowns: bool = True) -> References:
         """
-        Returns references object for specified certificate digest.
+        Get the references object for specified certificate digest.
 
-        :param str dgst: certificate digest
+        :param dgst: certificate digest
+        :param keep_unknowns: Whether to keep references to unknown certificate IDs
         :return References: References object
         """
         if not self._fitted:
             raise ValueError("Finder not yet fitted")
 
         def wrap(res):
-            return set(res) if res else None
+            if not res:
+                return None
+            # If we do not want the unknown references, filter them here.
+            if not keep_unknowns:
+                res = filter(lambda cert_id: cert_id in self.id_mapping, res)
+            return set(res)
 
         return References(
             wrap(self.dependencies[dgst].get("directly_referenced_by", None)),
@@ -148,11 +213,12 @@ class DependencyFinder:
             wrap(self.dependencies[dgst].get("indirectly_referencing", None)),
         )
 
-    def predict(self, dgst_list: List[str]) -> Dict[str, References]:
+    def predict(self, dgst_list: List[str], keep_unknowns: bool = True) -> Dict[str, References]:
         """
-        Returns references for a list of certificate digests.
+        Get the references for a list of certificate digests.
 
-        :param List[str] dgst_list: List of certificate digests.
+        :param dgst_list: List of certificate digests.
+        :param keep_unknowns: Whether to keep references to and from unknown certificate IDs
         :return Dict[str, References]: Dict with certificate hash and References object.
         """
         if not self._fitted:
@@ -160,6 +226,6 @@ class DependencyFinder:
         cert_references = {}
 
         for dgst in dgst_list:
-            cert_references[dgst] = self.predict_single_cert(dgst)
+            cert_references[dgst] = self.predict_single_cert(dgst, keep_unknowns=keep_unknowns)
 
         return cert_references