From 8a1fc0bd144d07e99801f276a6d7c1582dff8054 Mon Sep 17 00:00:00 2001 From: J08nY Date: Sun, 25 Sep 2022 18:16:29 +0200 Subject: Make DependencyFinder handle duplicates and unknowns. --- sec_certs/model/dependency_finder.py | 138 ++++++++++++++++++++++++++--------- 1 file 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 -- cgit v1.3.1