feat(general): Add image referencer rustworkx support

checkov/common/images/graph/image_referencer_manager.py

@@ -1,19 +1,18 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
+from checkov.common.typing import LibraryGraph
 
 if TYPE_CHECKING:
     from checkov.common.images.image_referencer import Image
-    from networkx import DiGraph
-    from igraph import Graph
 
 
 class GraphImageReferencerManager:
     __slots__ = ("graph_connector",)
 
-    def __init__(self, graph_connector: Union[Graph, DiGraph]) -> None:
+    def __init__(self, graph_connector: LibraryGraph) -> None:
         self.graph_connector = graph_connector
 
     @abstractmethod

checkov/common/images/graph/image_referencer_provider.py

@@ -2,15 +2,18 @@
 
 import itertools
 import os
+import typing
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Callable, Any, Mapping, Union, Generator
+from typing import TYPE_CHECKING, Callable, Any, Mapping, Generator
 
 from checkov.common.graph.graph_builder import CustomAttributes
 from checkov.common.images.image_referencer import Image
+from checkov.common.typing import LibraryGraph
 
 if TYPE_CHECKING:
     import networkx
     import igraph
+    import rustworkx
     from typing_extensions import TypeAlias
 
 _ExtractImagesCallableAlias: TypeAlias = Callable[["dict[str, Any]"], "list[str]"]
@@ -19,7 +22,7 @@
 class GraphImageReferencerProvider:
     __slots__ = ("graph_connector", "supported_resource_types", "graph_framework")
 
-    def __init__(self, graph_connector: Union[igraph.Graph, networkx.DiGraph],
+    def __init__(self, graph_connector: LibraryGraph,
                  supported_resource_types: dict[str, _ExtractImagesCallableAlias] | Mapping[
                      str, _ExtractImagesCallableAlias]):
         self.graph_connector = graph_connector
@@ -30,13 +33,17 @@ def __init__(self, graph_connector: Union[igraph.Graph, networkx.DiGraph],
     def extract_images_from_resources(self) -> list[Image]:
         pass
 
-    def extract_nodes(self) -> networkx.Graph | igraph.Graph | None:
+    def extract_nodes(self) -> LibraryGraph | None:
         if self.graph_framework == 'IGRAPH':
             return self.extract_nodes_igraph()
-        else:  # the default value of the graph framework is 'NETWORKX'
+        elif self.graph_framework == 'RUSTWORKX':
+            return self.extract_nodes_rustworkx()
+        else:
             return self.extract_nodes_networkx()
 
     def extract_nodes_networkx(self) -> networkx.Graph:
+        if typing.TYPE_CHECKING:
+            self.graph_connector = typing.cast(networkx.Graph, self.graph_connector)
         resource_nodes = [
             node
             for node, resource_type in self.graph_connector.nodes(data=CustomAttributes.RESOURCE_TYPE)
@@ -45,7 +52,18 @@ def extract_nodes_networkx(self) -> networkx.Graph:
 
         return self.graph_connector.subgraph(resource_nodes)
 
+    def extract_nodes_rustworkx(self) -> rustworkx.PyDiGraph[Any, Any]:
+        resource_nodes = [
+            index
+            for index, node in self.graph_connector.nodes()
+            if self.resource_type_pred(node, list(self.supported_resource_types))
+        ]
+
+        return self.graph_connector.subgraph(resource_nodes)
+
     def extract_nodes_igraph(self) -> igraph.Graph:
+        if typing.TYPE_CHECKING:
+            self.graph_connector = typing.cast(igraph.Graph, self.graph_connector)
         resource_nodes = [
             node
             for node, resource_type in itertools.zip_longest(self.graph_connector.vs['name'],
@@ -54,12 +72,16 @@ def extract_nodes_igraph(self) -> igraph.Graph:
         ]
         return self.graph_connector.subgraph(resource_nodes)
 
-    def extract_resource(self, supported_resources_graph: networkx.Graph | igraph.Graph) -> \
+    def extract_resource(self, supported_resources_graph: LibraryGraph) -> \
             Generator[dict[str, Any], dict[str, Any], dict[str, Any]]:
         def extract_resource_networkx(graph: networkx.Graph) -> Generator[dict[str, Any], None, None]:
             for _, resource in graph.nodes(data=True):
                 yield resource
 
+        def extract_resource_rustworkx(graph: rustworkx.PyDiGraph[Any, Any]) -> Generator[dict[str, Any], None, None]:
+            for _, resource in graph.nodes():
+                yield resource
+
         def extract_resource_igraph(graph: igraph.Graph) -> Generator[dict[str, Any], None, None]:
             for v in graph.vs:
                 resource = {
@@ -75,5 +97,11 @@ def extract_resource_igraph(graph: igraph.Graph) -> Generator[dict[str, Any], No
             graph_resource = extract_resource_networkx(supported_resources_graph)
         elif self.graph_framework == 'IGRAPH':
             graph_resource = extract_resource_igraph(supported_resources_graph)
+        elif self.graph_framework == 'RUSTWORKX':
+            graph_resource = extract_resource_rustworkx(supported_resources_graph)
 
         return graph_resource  # type: ignore
+
+    @staticmethod
+    def resource_type_pred(v: dict[str, Any], resource_types: list[str]) -> bool:
+        return not resource_types or ("resource_type" in v and v["resource_type"] in resource_types)

checkov/terraform/graph_builder/utils.py

@@ -8,11 +8,12 @@
 
 import igraph
 
+from checkov.common.typing import LibraryGraph
 from checkov.common.util.parser_utils import TERRAFORM_NESTED_MODULE_PATH_SEPARATOR_LENGTH, \
     TERRAFORM_NESTED_MODULE_INDEX_SEPARATOR
+from networkx import DiGraph
 
 if TYPE_CHECKING:
-    from networkx import DiGraph
     from checkov.terraform.graph_builder.graph_components.blocks import TerraformBlock
 
 from checkov.common.util.type_forcers import force_int
@@ -364,11 +365,25 @@ def get_file_path_to_referred_id_igraph(graph_object: igraph.Graph) -> dict[str,
     return file_path_to_module_id
 
 
-def setup_file_path_to_referred_id(graph_object: DiGraph | igraph.Graph) -> dict[str, str]:
+def get_file_path_to_referred_id_rustworkx(graph_object: DiGraph) -> dict[str, str]:
+    file_path_to_module_id = {}
+
+    modules = [node for index, node in graph_object.nodes() if
+               node.get(CustomAttributes.BLOCK_TYPE) == BlockType.MODULE]
+    for modules_data in modules:
+        for module_name, module_content in modules_data.get(CustomAttributes.CONFIG, {}).items():
+            for path in module_content.get("__resolved__", []):
+                file_path_to_module_id[path] = f"module.{module_name}"
+    return file_path_to_module_id
+
+
+def setup_file_path_to_referred_id(graph_object: LibraryGraph) -> dict[str, str]:
     if isinstance(graph_object, igraph.Graph):
         return get_file_path_to_referred_id_igraph(graph_object)
-    else:  # the default value of the graph framework is 'NETWORKX'
+    elif isinstance(graph_object, DiGraph):
         return get_file_path_to_referred_id_networkx(graph_object)
+    else:
+        return get_file_path_to_referred_id_rustworkx(graph_object)
 
 
 def get_attribute_is_leaf(vertex: TerraformBlock) -> Dict[str, bool]:

tests/ansible/checks/graph_checks/test_yaml_policies.py

@@ -1,13 +1,9 @@
-import os
 import warnings
 from pathlib import Path
 from typing import List
 from parameterized import parameterized_class
 
-from checkov.common.graph.db_connectors.networkx.networkx_db_connector import NetworkxConnector
-from checkov.common.graph.db_connectors.igraph.igraph_db_connector import IgraphConnector
-from checkov.common.graph.db_connectors.rustworkx.rustworkx_db_connector import RustworkxConnector
-from checkov.common.graph.db_connectors.utils import set_db_connector_by_graph_framework
+from tests.graph_utils.utils import set_db_connector_by_graph_framework, PARAMETERIZED_GRAPH_FRAMEWORKS
 
 from checkov.common.graph.graph_builder import CustomAttributes
 from checkov.common.models.enums import CheckResult
@@ -19,11 +15,7 @@
 from tests.common.graph.checks.test_yaml_policies_base import TestYamlPoliciesBase
 
 
-@parameterized_class([
-   {"graph_framework": "NETWORKX"},
-   {"graph_framework": "IGRAPH"},
-   {"graph_framework": "RUSTWORKX"}
-])
+@parameterized_class(PARAMETERIZED_GRAPH_FRAMEWORKS)
 class TestYamlPolicies(TestYamlPoliciesBase):
     def __init__(self, args):
         db_connector = set_db_connector_by_graph_framework(self.graph_framework)

tests/arm/graph_builder/checks/test_yaml_policies.py

@@ -3,11 +3,8 @@
 from typing import List
 
 from checkov.arm.graph_manager import ArmGraphManager
-from checkov.common.graph.db_connectors.networkx.networkx_db_connector import NetworkxConnector
 from parameterized import parameterized_class
-from checkov.common.graph.db_connectors.igraph.igraph_db_connector import IgraphConnector
-from checkov.common.graph.db_connectors.rustworkx.rustworkx_db_connector import RustworkxConnector
-from checkov.common.graph.db_connectors.utils import set_db_connector_by_graph_framework
+from tests.graph_utils.utils import set_db_connector_by_graph_framework, PARAMETERIZED_GRAPH_FRAMEWORKS
 
 from checkov.common.graph.graph_builder import CustomAttributes
 from checkov.common.models.enums import CheckResult
@@ -18,11 +15,7 @@
 
 
 @parameterized_class(
-    [
-        {"graph_framework": "NETWORKX"},
-        {"graph_framework": "IGRAPH"},
-        {"graph_framework": "RUSTWORKX"},
-    ]
+    PARAMETERIZED_GRAPH_FRAMEWORKS
 )
 class TestYamlPolicies(TestYamlPoliciesBase):
     def __init__(self, args):

tests/bicep/graph/checks/test_yaml_policies.py

@@ -3,11 +3,8 @@
 from typing import List
 
 from checkov.bicep.graph_manager import BicepGraphManager
-from checkov.common.graph.db_connectors.networkx.networkx_db_connector import NetworkxConnector
 from parameterized import parameterized_class
-from checkov.common.graph.db_connectors.igraph.igraph_db_connector import IgraphConnector
-from checkov.common.graph.db_connectors.rustworkx.rustworkx_db_connector import RustworkxConnector
-from checkov.common.graph.db_connectors.utils import set_db_connector_by_graph_framework
+from tests.graph_utils.utils import set_db_connector_by_graph_framework, PARAMETERIZED_GRAPH_FRAMEWORKS
 
 from checkov.common.graph.graph_builder import CustomAttributes
 from checkov.common.models.enums import CheckResult
@@ -16,11 +13,7 @@
 from checkov.common.bridgecrew.check_type import CheckType
 from tests.common.graph.checks.test_yaml_policies_base import TestYamlPoliciesBase
 
-@parameterized_class([
-   {"graph_framework": "NETWORKX"},
-   {"graph_framework": "IGRAPH"},
-   {"graph_framework": "RUSTWORKX"}
-])
+@parameterized_class(PARAMETERIZED_GRAPH_FRAMEWORKS)
 class TestYamlPolicies(TestYamlPoliciesBase):
     def __init__(self, args):
         db_connector = set_db_connector_by_graph_framework(self.graph_framework)

tests/bicep/image_referencer/provider/test_azure.py

@@ -1,17 +1,18 @@
-import os
-
 from unittest import mock
 
 import igraph
 import pytest
 from networkx import DiGraph
+from rustworkx import PyDiGraph
 
 from checkov.common.graph.graph_builder import CustomAttributes
 from checkov.common.images.image_referencer import Image
 from checkov.terraform.image_referencer.provider.azure import AzureTerraformProvider
+from tests.graph_utils.utils import GRAPH_FRAMEWORKS, set_graph_by_graph_framework, \
+    add_vertices_to_graph_by_graph_framework
 
 
-@pytest.mark.parametrize("graph_framework", ['NETWORKX', 'IGRAPH'])
+@pytest.mark.parametrize("graph_framework", GRAPH_FRAMEWORKS)
 def extract_images_from_resources(graph_framework):
     # given
     resource = {
@@ -33,18 +34,8 @@ def extract_images_from_resources(graph_framework):
             },
             "resource_type": "Microsoft.Batch/batchAccounts/pools",
         }
-    if graph_framework == 'IGRAPH':
-        graph = igraph.Graph()
-        graph.add_vertex(
-            name='1',
-            block_type_='resource',
-            resource_type=resource[
-                CustomAttributes.RESOURCE_TYPE] if CustomAttributes.RESOURCE_TYPE in resource else None,
-            attr=resource,
-        )
-    else:
-        graph = DiGraph()
-        graph.add_node(1, **resource)
+    graph = set_graph_by_graph_framework(graph_framework)
+    add_vertices_to_graph_by_graph_framework(graph_framework, resource, graph)
 
     # when
     azure_provider = AzureTerraformProvider(graph_connector=graph)
@@ -57,7 +48,7 @@ def extract_images_from_resources(graph_framework):
     ]
 
 
-@pytest.mark.parametrize('graph_framework', ["NETWORKX", "IGRAPH"])
+@pytest.mark.parametrize('graph_framework', GRAPH_FRAMEWORKS)
 def test_extract_images_from_resources_with_no_image(graph_framework):
     # given
     resource = {
@@ -79,18 +70,8 @@ def test_extract_images_from_resources_with_no_image(graph_framework):
         },
         "resource_type": "Microsoft.Batch/batchAccounts/pools",
     }
-    if graph_framework == 'IGRAPH':
-        graph = igraph.Graph()
-        graph.add_vertex(
-            name='1',
-            block_type_='resource',
-            resource_type=resource[
-                CustomAttributes.RESOURCE_TYPE] if CustomAttributes.RESOURCE_TYPE in resource else None,
-            attr=resource,
-        )
-    else:
-        graph = DiGraph()
-        graph.add_node(1, **resource)
+    graph = set_graph_by_graph_framework(graph_framework)
+    add_vertices_to_graph_by_graph_framework(graph_framework, resource, graph)
 
     # when
     with mock.patch.dict('os.environ', {'CHECKOV_GRAPH_FRAMEWORK': graph_framework}):

tests/bicep/image_referencer/test_manager.py

@@ -1,17 +1,14 @@
-import unittest
 from unittest import mock
 
-import igraph
 import pytest
-from networkx import DiGraph
-from parameterized import parameterized_class
 
 from checkov.bicep.image_referencer.manager import BicepImageReferencerManager
-from checkov.common.graph.graph_builder import CustomAttributes
 from checkov.common.images.image_referencer import Image
+from tests.graph_utils.utils import GRAPH_FRAMEWORKS, set_graph_by_graph_framework, \
+    add_vertices_to_graph_by_graph_framework
 
 
-@pytest.mark.parametrize("graph_framework", ['NETWORKX', 'IGRAPH'])
+@pytest.mark.parametrize("graph_framework", GRAPH_FRAMEWORKS)
 def test_extract_images_from_resources(graph_framework):
     # given
     resource = {
@@ -33,18 +30,9 @@ def test_extract_images_from_resources(graph_framework):
         },
         "resource_type": "Microsoft.Batch/batchAccounts/pools",
     }
-    if graph_framework == 'NETWORKX':
-        graph = DiGraph()
-        graph.add_node(1, **resource)
-    else:
-        graph = igraph.Graph()
-        graph.add_vertex(
-            name='1',
-            block_type_='resource',
-            resource_type=resource[CustomAttributes.RESOURCE_TYPE] if CustomAttributes.RESOURCE_TYPE in
-                                                                      resource else None,
-            attr=resource,
-        )
+    graph = set_graph_by_graph_framework(graph_framework)
+    add_vertices_to_graph_by_graph_framework(graph_framework, resource, graph)
+
 
     # when
     with mock.patch.dict('os.environ', {'CHECKOV_GRAPH_FRAMEWORK': graph_framework}):