[azure] [feat] Add new relationships between network resources (#1838)

plugins/azure/resoto_plugin_azure/azure_client.py

@@ -133,30 +133,45 @@ def _update_or_delete_tag(self, tag_name: str, tag_value: str, resource_id: str,
 
     # noinspection PyProtectedMember
     def _call(self, spec: AzureApiSpec, **kwargs: Any) -> List[Json]:
-        _SERIALIZER = Serializer()
+        ser = Serializer()
 
         error_map = {
             401: ClientAuthenticationError,
             404: ResourceNotFoundError,
         }
 
+        # Construct lookup map used to fill query and path parameters
+        lookup_map = {"subscriptionId": self.subscription_id, "location": self.location, **kwargs}
+
         # Construct headers
-        headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
-        headers["Accept"] = _SERIALIZER.header("accept", headers.pop("Accept", "application/json"), "str")  # type: ignore # noqa: E501
+        headers = case_insensitive_dict()
+        headers["Accept"] = ser.header("accept", headers.pop("Accept", "application/json"), "str")  # type: ignore # noqa: E501
+
+        # Construct path map
+        path_map = case_insensitive_dict()
+        for param in spec.path_parameters:
+            if lookup_map.get(param, None) is not None:
+                path_map[param] = lookup_map[param]
+            else:
+                raise ValueError(f"Param {param} in lookup_map does not found")
 
         # Construct parameters
-        params = case_insensitive_dict(kwargs.pop("params", {}) or {})
-        params["api-version"] = _SERIALIZER.query("api_version", spec.version, "str")  # type: ignore
+        params = case_insensitive_dict()
+        params["api-version"] = ser.query("api-version", spec.version, "str")  # type: ignore
+        for param in spec.query_parameters:
+            if param not in params:
+                if lookup_map.get(param, None) is not None:
+                    params[param] = ser.query(param, lookup_map[param], "str")  # type: ignore # noqa: E501
+                else:
+                    raise ValueError(f"Param {param} in lookup_map does not found")
 
         # Construct url
-        path = spec.path.format_map({"subscriptionId": self.subscription_id, "location": self.location, **params})
+        path = spec.path.format_map(path_map)
         url = self.client._client.format_url(path)  # pylint: disable=protected-access
 
         # Construct and send request
-        request = HttpRequest(method="GET", url=url, params=params, headers=headers, **kwargs)
-        pipeline_response: PipelineResponse = self.client._client._pipeline.run(  # type: ignore
-            request, stream=False, **kwargs
-        )
+        request = HttpRequest(method="GET", url=url, params=params, headers=headers)
+        pipeline_response: PipelineResponse = self.client._client._pipeline.run(request, stream=False)  # type: ignore
         response = pipeline_response.http_response
 
         # Handle error responses
@@ -165,7 +180,6 @@ def _call(self, spec: AzureApiSpec, **kwargs: Any) -> List[Json]:
             raise HttpResponseError(response=response, error_format=ARMErrorFormat)
 
         # Parse json content
-        # TODO: handle pagination
         js: Union[Json, List[Json]] = response.json()
         if spec.access_path and isinstance(js, dict):
             js = js[spec.access_path]

plugins/azure/resoto_plugin_azure/collector.py

@@ -65,7 +65,6 @@ def collect(self) -> None:
             # collect all regional resources
             for location in locations:
                 self.collect_resource_list(location.safe_name, builder.with_location(location), regional_resources)
-
             # wait for all work to finish
             queue.wait_for_submitted_work()
             # connect nodes

plugins/azure/resoto_plugin_azure/resource/base.py

@@ -4,7 +4,7 @@
 from concurrent.futures import Future
 from datetime import datetime
 from threading import Lock
-from typing import Any, ClassVar, Dict, Optional, TypeVar, List, Type, Callable, cast
+from typing import Any, ClassVar, Dict, Optional, TypeVar, List, Tuple, Type, Callable, Union, cast
 
 from attr import define, field
 from azure.core.utils import CaseInsensitiveDict
@@ -45,15 +45,48 @@ class AzureResource(BaseResource):
     api_spec: ClassVar[Optional[AzureApiSpec]] = None
 
     def resource_subscription_id(self) -> str:
+        return self.extract_part("subscriptionId")
+
+    def resource_group_name(self) -> str:
+        return self.extract_part("resourceGroupName")
+
+    def extract_part(self, part: str) -> str:
         """
-        Extracts {subscriptionId} value from a resource ID.
+        Extracts a specific part from a resource ID.
+
+        The function takes a resource ID and a specified part to extract, such as 'subscriptionId'
+        or 'resourceGroupName'. The resource ID is expected to follow the Azure Resource Manager
+        path format.
+
+        Example:
+        For the resource ID "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/...",
+        calling extract_part("subscriptionId") would return the value within the curly braces,
+        representing the subscription ID.
 
-        e.g. /subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/...
+        Parameters:
+        - part (str): The part to extract from the resource ID.
 
         Returns:
-        str: The extracted subscription ID.
+        str: The extracted part of the resource ID.
         """
-        return self.id.split("/")[2]
+        id_parts = self.id.split("/")
+
+        if part == "subscriptionId":
+            if "subscriptions" not in id_parts:
+                raise ValueError(f"Id {self.id} does not have any subscriptionId info")
+            if index := id_parts.index("subscriptions"):
+                return id_parts[index + 1]
+            return ""
+
+        elif part == "resourceGroupName":
+            if "resourceGroups" not in id_parts:
+                raise ValueError(f"Id {self.id} does not have any resourceGroupName info")
+            if index := id_parts.index("resourceGroups"):
+                return id_parts[index + 1]
+            return ""
+
+        else:
+            raise ValueError(f"Value {part} does not have any cases to match")
 
     def delete(self, graph: Graph) -> bool:
         """
@@ -101,6 +134,35 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         # Default behavior: add resource to the namespace
         pass
 
+    def fetch_resources(
+        self,
+        builder: GraphBuilder,
+        service: str,
+        api_version: str,
+        path: str,
+        path_parameters: List[str],
+        query_parameters: List[str],
+        compared_property: Callable[[Json], Union[List[str], str]],
+        binding_property: Callable[[Json], Union[List[str], str]],
+    ) -> List[Tuple[Union[str, List[str]], Union[str, List[str]]]]:
+        """
+        Fetch additional resources from the Azure API for further connection using the connect_in_graph method.
+
+        Returns:
+        List[Tuple[Union[str, List[str]], str]]: A list of tuples containing information to compare and connect the retrieved resources.
+        """
+        resources_api_spec = AzureApiSpec(
+            service=service,
+            version=api_version,
+            path=path,
+            path_parameters=path_parameters,
+            query_parameters=query_parameters,
+            access_path="value",
+            expect_array=True,
+        )
+
+        return [(compared_property(r), binding_property(r)) for r in builder.client.list(resources_api_spec)]
+
     @classmethod
     def collect_resources(
         cls: Type[AzureResourceType], builder: GraphBuilder, **kwargs: Any

plugins/azure/resoto_plugin_azure/resource/compute.py

@@ -14,6 +14,12 @@
     AzurePrincipalidClientid,
     AzurePrivateLinkServiceConnectionState,
 )
+from resoto_plugin_azure.resource.network import (
+    AzureNetworkSecurityGroup,
+    AzureSubnet,
+    AzureNetworkInterface,
+    AzureLoadBalancer,
+)
 from resotolib.json_bender import Bender, S, Bend, MapEnum, ForallBend, K, F
 from resotolib.types import Json
 from resotolib.baseresources import (
@@ -353,7 +359,7 @@ class AzureComputeOperationValue(AzureResource):
         expect_array=True,
     )
     mapping: ClassVar[Dict[str, Bender]] = {
-        "id": K(None),
+        "id": S("name"),
         "tags": S("tags", default={}),
         "name": S("name"),
         "ctime": K(None),
@@ -714,7 +720,7 @@ class AzureDisk(AzureResource, BaseVolume):
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if disk_id := self.id:
-            builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureDiskAccess, id=disk_id)
+            builder.add_edge(self, edge_type=EdgeType.default, reverse=True, clazz=AzureDiskAccess, id=disk_id)
         if (disk_encryption := self.disk_encryption) and (disk_en_set_id := disk_encryption.disk_encryption_set_id):
             builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureDiskEncryptionSet, id=disk_en_set_id)
 
@@ -1065,6 +1071,9 @@ class AzureProximityPlacementGroup(AzureResource):
         access_path="value",
         expect_array=True,
     )
+    reference_kinds: ClassVar[ModelReference] = {
+        "successors": {"default": ["azure_virtual_machine_scale_set"]},
+    }
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("id"),
         "tags": S("tags", default={}),
@@ -1089,6 +1098,12 @@ class AzureProximityPlacementGroup(AzureResource):
     virtual_machine_scale_sets: Optional[List[AzureSubResourceWithColocationStatus]] = field(default=None, metadata={'description': 'A list of references to all virtual machine scale sets in the proximity placement group.'})  # fmt: skip
     virtual_machines_status: Optional[List[AzureSubResourceWithColocationStatus]] = field(default=None, metadata={'description': 'A list of references to all virtual machines in the proximity placement group.'})  # fmt: skip
 
+    def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
+        if vmsss := self.virtual_machine_scale_sets:
+            for vmss in vmsss:
+                if vmss_id := vmss.id:
+                    builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureVirtualMachineScaleSet, id=vmss_id)
+
 
 @define(eq=False, slots=False)
 class AzureResourceSkuCapacity:
@@ -1843,7 +1858,7 @@ class AzureSnapshot(AzureResource, BaseSnapshot):
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if (disk_data := self.creation_data) and (disk_id := disk_data.source_resource_id):
-            builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureDisk, id=disk_id)
+            builder.add_edge(self, edge_type=EdgeType.default, reverse=True, clazz=AzureDisk, id=disk_id)
 
 
 @define(eq=False, slots=False)
@@ -2505,7 +2520,21 @@ class AzureVirtualMachine(AzureResource, BaseInstance):
         expect_array=True,
     )
     reference_kinds: ClassVar[ModelReference] = {
-        "successors": {"default": ["azure_proximity_placement_group", "azure_image", "azure_disk"]},
+        "predecessors": {
+            "default": [
+                "azure_proximity_placement_group",
+                "azure_network_security_group",
+                "azure_subnet",
+                "azure_load_balancer",
+            ]
+        },
+        "successors": {
+            "default": [
+                "azure_image",
+                "azure_disk",
+                "azure_network_interface",
+            ]
+        },
     }
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("id"),
@@ -2587,7 +2616,11 @@ class AzureVirtualMachine(AzureResource, BaseInstance):
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if placement_group_id := self.proximity_placement_group:
             builder.add_edge(
-                self, edge_type=EdgeType.default, clazz=AzureProximityPlacementGroup, id=placement_group_id
+                self,
+                edge_type=EdgeType.default,
+                reverse=True,
+                clazz=AzureProximityPlacementGroup,
+                id=placement_group_id,
             )
 
         if (
@@ -2605,6 +2638,35 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         ):
             builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureDisk, id=managed_disk_id)
 
+        if (vm_network_profile := self.virtual_machine_network_profile) and (
+            ni_cofigurations := vm_network_profile.network_interface_configurations
+        ):
+            for ni_configuration in ni_cofigurations:
+                if nsg_id := ni_configuration.network_security_group:
+                    builder.add_edge(
+                        self, edge_type=EdgeType.default, reverse=True, clazz=AzureNetworkSecurityGroup, id=nsg_id
+                    )
+                if ip_configurations := ni_configuration.ip_configurations:
+                    for ip_configuration in ip_configurations:
+                        if subnet_id := ip_configuration.subnet:
+                            builder.add_edge(
+                                self, edge_type=EdgeType.default, reverse=True, clazz=AzureSubnet, id=subnet_id
+                            )
+                        if lbbap_ids := ip_configuration.load_balancer_backend_address_pools:
+                            for lbbap_id in lbbap_ids:
+                                # take only id of load balancer
+                                lbbap_id = "/".join(lbbap_id.split("/")[:-2])
+                                builder.add_edge(
+                                    self, edge_type=EdgeType.default, reverse=True, clazz=AzureLoadBalancer, id=lbbap_id
+                                )
+
+        if (vm_network_profile := self.virtual_machine_network_profile) and (
+            network_interfaces := vm_network_profile.network_interfaces
+        ):
+            for network_interface in network_interfaces:
+                if ni_id := network_interface.id:
+                    builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureNetworkInterface, id=ni_id)
+
 
 @define(eq=False, slots=False)
 class AzureRollingUpgradePolicy:
@@ -3054,6 +3116,9 @@ class AzureVirtualMachineScaleSet(AzureResource, BaseAutoScalingGroup):
         access_path="value",
         expect_array=True,
     )
+    reference_kinds: ClassVar[ModelReference] = {
+        "predecessors": {"default": ["azure_load_balancer"]},
+    }
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("id"),
         "tags": S("tags", default={}),
@@ -3112,6 +3177,26 @@ class AzureVirtualMachineScaleSet(AzureResource, BaseAutoScalingGroup):
     virtual_machine_profile: Optional[AzureVirtualMachineScaleSetVMProfile] = field(default=None, metadata={'description': 'Describes a virtual machine scale set virtual machine profile.'})  # fmt: skip
     zone_balance: Optional[bool] = field(default=None, metadata={'description': 'Whether to force strictly even virtual machine distribution cross x-zones in case there is zone outage. Zonebalance property can only be set if the zones property of the scale set contains more than one zone. If there are no zones or only one zone specified, then zonebalance property should not be set.'})  # fmt: skip
 
+    def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
+        if (
+            (vm_profile := self.virtual_machine_profile)
+            and (net_profile := vm_profile.network_profile)
+            and (net_i_configs := net_profile.network_interface_configurations)
+        ):
+            for net_i_config in net_i_configs:
+                if ip_configs := net_i_config.ip_configurations:
+                    for ip_config in ip_configs:
+                        if baps := ip_config.load_balancer_backend_address_pools:
+                            for bap in baps:
+                                if bap_id := bap:
+                                    builder.add_edge(
+                                        self,
+                                        edge_type=EdgeType.default,
+                                        reverse=True,
+                                        clazz=AzureLoadBalancer,
+                                        id=bap_id,
+                                    )
+
 
 @define(eq=False, slots=False)
 class AzureVirtualMachineSize(AzureResource, BaseInstanceType):
@@ -3126,7 +3211,7 @@ class AzureVirtualMachineSize(AzureResource, BaseInstanceType):
         expect_array=True,
     )
     mapping: ClassVar[Dict[str, Bender]] = {
-        "id": S("id"),
+        "id": S("name"),
         "tags": S("tags", default={}),
         "name": S("name"),
         "ctime": K(None),