Merge pull request #47 from RedHatProductSecurity/PSDEVOPS-4471

PSDEVOPS-4471

Author: jasinner

src/trustshell/product_definitions.py

@@ -210,27 +210,42 @@ def _clean_cpe(cpe: str) -> str:
         # Remove trailing ':' characters
         return cleaned_cpe.rstrip(":")
 
-    def extend_with_product_mappings(self, ancestor_trees: list[Node]) -> list[Node]:
-        """Create a new list of results with any matching streams or module as ancestors"""
+    def extend_with_product_mappings(
+        self, ancestor_trees: list[Node], keep_cpes: bool = False
+    ) -> None:
+        """Update the ancestor_trees with any matching streams or module as descendants
+
+        Args:
+            ancestor_trees: List of Node trees to extend with product mappings
+            keep_cpes: If False, replace CPE leaf nodes with product streams. If True, keep CPE nodes as parents of streams.
+        """
         if not self.product_trees:
             # ProdDefs service is unavailable, don't attempt any product mapping
-            return ancestor_trees
+            return None
+
         for tree in ancestor_trees:
             for leaf in tree.leaves:
                 cleaned_leaf_name = self._clean_cpe(leaf.name)
-                leaf_with_products = self._check_streams(leaf, cleaned_leaf_name)
+                leaf_with_products = self._check_streams(
+                    leaf, cleaned_leaf_name, keep_cpes
+                )
                 if not leaf_with_products:
-                    leaf_with_products = self._check_modules(leaf, cleaned_leaf_name)
+                    leaf_with_products = self._check_modules(
+                        leaf, cleaned_leaf_name, keep_cpes
+                    )
                 if not leaf_with_products:
                     console.print(
                         f"Warning, didn't find any products matching {cleaned_leaf_name}",
                         style="warning",
                     )
-                # the tree is modified as a side effect of the checking streams|modules
-                # therefore we do not need to explicitly store and return it elsewhere
-        return ancestor_trees
-
-    def _check_streams(self, leaf: Node, cpe: str) -> list[Node]:
+                else:
+                    # When keep_cpes=False, we need to remove the CPE leaf from the tree
+                    # since it's been replaced by the product streams
+                    if not keep_cpes:
+                        # Remove the CPE leaf node from its parent
+                        leaf.parent = None
+
+    def _check_streams(self, leaf: Node, cpe: str, keep_cpes: bool) -> list[Node]:
         """Check if cpe matches exactly to any ProductStreams, if it does add the CPE as a parent
         of the stream. If more than one stream matches, create copies of the stream and leaf"""
         if cpe not in self.stream_nodes_by_cpe:
@@ -240,18 +255,30 @@ def _check_streams(self, leaf: Node, cpe: str) -> list[Node]:
         # the original stream_nodes_by_cpe map which should be preserved incase we encounter the
         # same CPE twice
         copy_of_stream_nodes = copy.deepcopy(stream_nodes)
-        return self._duplicate_leaves_and_set_parents(leaf, copy_of_stream_nodes)
+        return self._duplicate_leaves_and_set_parents(
+            leaf, copy_of_stream_nodes, keep_cpes
+        )
 
-    def _check_modules(self, leaf: Node, cpe: str) -> list[Node]:
+    def _check_modules(self, leaf: Node, cpe: str, keep_cpes: bool) -> list[Node]:
         """Check if the cpe matches any ProductModule"""
         module_nodes = self.match_module_pattern(cpe)
-        return self._duplicate_leaves_and_set_parents(leaf, module_nodes)
+        return self._duplicate_leaves_and_set_parents(leaf, module_nodes, keep_cpes)
 
     def _duplicate_leaves_and_set_parents(
-        self, leaf: Node, product_nodes: list[Any]
+        self, leaf: Node, product_nodes: list[Any], keep_cpes: bool
     ) -> list[Node]:
         """Convert product modules to their parent streams and attach all streams as children of the leaf.
-        Deduplicates streams to avoid multiple identical children. Returns the leaf in a list."""
+        Deduplicates streams to avoid multiple identical children.
+
+        Args:
+            leaf: The leaf node to process
+            product_nodes: List of product nodes (modules or streams)
+            keep_cpes: If False, replace the leaf with streams in the tree. If True, set streams as children of the leaf.
+
+        Returns:
+            If keep_cpes=True: Returns the leaf in a list.
+            If keep_cpes=False: Returns the list of unique streams that replaced the leaf.
+        """
         if not product_nodes:
             return []
 
@@ -274,10 +301,17 @@ def _duplicate_leaves_and_set_parents(
                 unique_streams.append(stream)
                 seen.add(stream)
 
-        for stream in unique_streams:
-            stream.parent = leaf
-
-        return [leaf]
+        if keep_cpes:
+            # Original behavior: set streams as children of the leaf
+            for stream in unique_streams:
+                stream.parent = leaf
+            return [leaf]
+        else:
+            # New behavior: replace the leaf with the streams in the tree
+            # Each stream takes the place of the leaf in the tree hierarchy
+            for stream in unique_streams:
+                stream.parent = leaf.parent
+            return unique_streams
 
     def _add_ancestor(self, leaf: Node, product: Any) -> None:
         if product.parent:

src/trustshell/products.py

@@ -72,6 +72,7 @@ def prime_cache(check: bool, debug: bool) -> None:
     is_eager=True,
 )
 @click.option("--latest", "-l", is_flag=True, default=True)
+@click.option("--cpes", "-c", is_flag=True, default=False)
 @click.option("--flaw", "-f", help="OSIDB flaw uuid or CVE")
 @click.option(
     "--replace",
@@ -85,7 +86,9 @@ def prime_cache(check: bool, debug: bool) -> None:
     "purl",
     type=click.STRING,
 )
-def search(purl: str, flaw: str, replace: bool, debug: bool, latest: bool) -> None:
+def search(
+    purl: str, flaw: str, replace: bool, debug: bool, latest: bool, cpes: bool
+) -> None:
     """Relate a purl to products in Trustify"""
     if not debug:
         config_logging(level="INFO")
@@ -104,9 +107,15 @@ def search(purl: str, flaw: str, replace: bool, debug: bool, latest: bool) -> No
         return
 
     prod_defs = ProdDefs()
-    ancestor_trees = prod_defs.extend_with_product_mappings(ancestor_trees)
+    prod_defs.extend_with_product_mappings(ancestor_trees, keep_cpes=cpes)
 
+    seen_trees = set()
     for tree in ancestor_trees:
+        _remove_duplicate_branches(tree)
+        tree_signature = _get_branch_signature(tree.root)
+        if tree_signature in seen_trees:
+            continue
+        seen_trees.add(tree_signature)
         render_tree(tree.root)
 
     if not flaw:
@@ -141,27 +150,31 @@ def extract_affects(ancestor_trees: list[Node]) -> set[tuple[str, str]]:
             raise ValueError(f"More than one ProductModule found in {tree.root.name}")
         for ps_module_node in ps_module_nodes:
             for ancestor in ps_module_node.ancestors:
-                if ancestor.name.startswith("cpe:/"):
-                    if ancestor.parent:
-                        purl = PackageURL.from_string(ancestor.parent.name)
-                        if purl.type == "oci" and "tag" in purl.qualifiers:
-                            purl.qualifiers.pop("tag")
-                        elif (
-                            purl.type == "maven"
-                            or purl.type == "generic"
-                            and PackageURL.from_string(ps_module_node.root.name).type
-                            == "maven"
-                        ):
-                            # If it's a maven type or a generic one based on maven,
-                            # we set the purl to root
-                            purl = PackageURL.from_string(ps_module_node.root.name)
-                        purl_sans_version_obj = purl_sans_version(purl)
-                        affects.add(
-                            (
-                                ps_module_node.name,
-                                purl_sans_version_obj.to_string(),
-                            )
-                        )
+                # Find the root component
+                if ancestor.name.startswith("pkg:"):
+                    purl = PackageURL.from_string(ancestor.name)
+                else:
+                    continue
+                # Clean up the purl ready for use in affects
+                if purl.type == "oci" and "tag" in purl.qualifiers:
+                    purl.qualifiers.pop("tag")
+                elif (
+                    purl.type == "maven"
+                    or purl.type == "generic"
+                    and PackageURL.from_string(ps_module_node.root.name).type == "maven"
+                ):
+                    # If it's a maven type or a generic one based on maven,
+                    # we set the purl to root
+                    purl = PackageURL.from_string(ps_module_node.root.name)
+                else:
+                    purl = purl_sans_version(purl)
+
+                affects.add(
+                    (
+                        ps_module_node.name,
+                        purl.to_string(),
+                    )
+                )
     return affects
 
 
@@ -353,14 +366,15 @@ def _trees_with_cpes(ancestor_data: dict[str, Any]) -> list[Node]:
         if tree.name.startswith("pkg:rpm/"):
             if container_in_tree(tree):
                 continue
+        _remove_non_cpe_branches(tree)
         if not _has_cpe_node(tree):
             for leaf in tree.leaves:
                 logger.debug(
                     f"Found result {tree.name} with ancestor: {leaf.name} but no CPE parent"
                 )
         else:
             trees_with_cpes.append(tree)
-    return [_remove_non_cpe_branches(tree) for tree in trees_with_cpes]
+    return trees_with_cpes
 
 
 def container_in_tree(root: Node) -> bool: