cartography 0.104.0rc2__py3-none-any.whl → 0.123.0__py3-none-any.whl

cartography/intel/azure/compute.py

@@ -1,12 +1,18 @@
 import logging
 from typing import Dict
 from typing import List
+from typing import Tuple
 
 import neo4j
 from azure.core.exceptions import HttpResponseError
 from azure.mgmt.compute import ComputeManagementClient
 
-from cartography.util import run_cleanup_job
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.azure.vm.datadisk import AzureDataDiskSchema
+from cartography.models.azure.vm.disk import AzureDiskSchema
+from cartography.models.azure.vm.snapshot import AzureSnapshotSchema
+from cartography.models.azure.vm.virtualmachine import AzureVirtualMachineSchema
 from cartography.util import timeit
 
 from .util.credentials import Credentials
@@ -44,70 +50,27 @@ def load_vms(
     vm_list: List[Dict],
     update_tag: int,
 ) -> None:
-    ingest_vm = """
-    UNWIND $vms AS vm
-    MERGE (v:AzureVirtualMachine{id: vm.id})
-    ON CREATE SET v.firstseen = timestamp(),
-    v.type = vm.type, v.location = vm.location,
-    v.resourcegroup = vm.resource_group
-    SET v.lastupdated = $update_tag, v.name = vm.name,
-    v.plan = vm.plan.product, v.size = vm.hardware_profile.vm_size,
-    v.license_type=vm.license_type, v.computer_name=vm.os_profile.computer_name,
-    v.identity_type=vm.identity.type, v.zones=vm.zones,
-    v.ultra_ssd_enabled=vm.additional_capabilities.ultra_ssd_enabled,
-    v.priority=vm.priority, v.eviction_policy=vm.eviction_policy
-    WITH v
-    MATCH (owner:AzureSubscription{id: $SUBSCRIPTION_ID})
-    MERGE (owner)-[r:RESOURCE]->(v)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-    """
-
-    neo4j_session.run(
-        ingest_vm,
-        vms=vm_list,
-        SUBSCRIPTION_ID=subscription_id,
-        update_tag=update_tag,
+    load(
+        neo4j_session,
+        AzureVirtualMachineSchema(),
+        vm_list,
+        lastupdated=update_tag,
+        AZURE_SUBSCRIPTION_ID=subscription_id,
     )
 
-    for vm in vm_list:
-        if vm.get("storage_profile", {}).get("data_disks"):
-            load_vm_data_disks(
-                neo4j_session,
-                vm["id"],
-                vm["storage_profile"]["data_disks"],
-                update_tag,
-            )
-
 
 def load_vm_data_disks(
     neo4j_session: neo4j.Session,
-    vm_id: str,
+    subscription_id: str,
     data_disks: List[Dict],
     update_tag: int,
 ) -> None:
-    ingest_data_disk = """
-    UNWIND $disks AS disk
-    MERGE (d:AzureDataDisk{id: disk.managed_disk.id})
-    ON CREATE SET d.firstseen = timestamp(), d.lun = disk.lun
-    SET d.lastupdated = $update_tag, d.name = disk.name,
-    d.vhd = disk.vhd.uri, d.image = disk.image.uri,
-    d.size = disk.disk_size_gb, d.caching = disk.caching,
-    d.createoption = disk.create_option, d.write_accelerator_enabled=disk.write_accelerator_enabled,
-    d.managed_disk_storage_type=disk.managed_disk.storage_account_type
-    WITH d
-    MATCH (owner:AzureVirtualMachine{id: $VM_ID})
-    MERGE (owner)-[r:ATTACHED_TO]->(d)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag
-    """
-
-    # for disk in data_disks:
-    neo4j_session.run(
-        ingest_data_disk,
-        disks=data_disks,
-        VM_ID=vm_id,
-        update_tag=update_tag,
+    load(
+        neo4j_session,
+        AzureDataDiskSchema(),
+        data_disks,
+        lastupdated=update_tag,
+        AZURE_SUBSCRIPTION_ID=subscription_id,
     )
 
 
@@ -115,10 +78,8 @@ def cleanup_virtual_machine(
     neo4j_session: neo4j.Session,
     common_job_parameters: Dict,
 ) -> None:
-    run_cleanup_job(
-        "azure_import_virtual_machines_cleanup.json",
+    GraphJob.from_node_schema(AzureVirtualMachineSchema(), common_job_parameters).run(
         neo4j_session,
-        common_job_parameters,
     )
 
 
@@ -144,37 +105,18 @@ def load_disks(
     disk_list: List[Dict],
     update_tag: int,
 ) -> None:
-    ingest_disks = """
-    UNWIND $disks AS disk
-    MERGE (d:AzureDisk{id: disk.id})
-    ON CREATE SET d.firstseen = timestamp(),
-    d.type = disk.type, d.location = disk.location,
-    d.resourcegroup = disk.resource_group
-    SET d.lastupdated = $update_tag, d.name = disk.name,
-    d.createoption = disk.creation_data.create_option, d.disksizegb = disk.disk_size_gb,
-    d.encryption = disk.encryption_settings_collection.enabled, d.maxshares = disk.max_shares,
-    d.network_access_policy = disk.network_access_policy,
-    d.ostype = disk.os_type, d.tier = disk.tier,
-    d.sku = disk.sku.name, d.zones = disk.zones
-    WITH d
-    MATCH (owner:AzureSubscription{id: $SUBSCRIPTION_ID})
-    MERGE (owner)-[r:RESOURCE]->(d)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag"""
-
-    neo4j_session.run(
-        ingest_disks,
-        disks=disk_list,
-        SUBSCRIPTION_ID=subscription_id,
-        update_tag=update_tag,
+    load(
+        neo4j_session,
+        AzureDiskSchema(),
+        disk_list,
+        lastupdated=update_tag,
+        AZURE_SUBSCRIPTION_ID=subscription_id,
     )
 
 
 def cleanup_disks(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job(
-        "azure_import_disks_cleanup.json",
+    GraphJob.from_node_schema(AzureDiskSchema(), common_job_parameters).run(
         neo4j_session,
-        common_job_parameters,
     )
 
 
@@ -200,39 +142,37 @@ def load_snapshots(
     snapshots: List[Dict],
     update_tag: int,
 ) -> None:
-    ingest_snapshots = """
-    UNWIND $snapshots as snapshot
-    MERGE (s:AzureSnapshot{id: snapshot.id})
-    ON CREATE SET s.firstseen = timestamp(),
-    s.resourcegroup = snapshot.resource_group,
-    s.type = snapshot.type, s.location = snapshot.location
-    SET s.lastupdated = $update_tag, s.name = snapshot.name,
-    s.createoption = snapshot.creation_data.create_option, s.disksizegb = snapshot.disk_size_gb,
-    s.encryption = snapshot.encryption_settings_collection.enabled, s.incremental = snapshot.incremental,
-    s.network_access_policy = snapshot.network_access_policy, s.ostype = snapshot.os_type,
-    s.tier = snapshot.tier, s.sku = snapshot.sku.name
-    WITH s
-    MATCH (owner:AzureSubscription{id: $SUBSCRIPTION_ID})
-    MERGE (owner)-[r:RESOURCE]->(s)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $update_tag"""
-
-    neo4j_session.run(
-        ingest_snapshots,
-        snapshots=snapshots,
-        SUBSCRIPTION_ID=subscription_id,
-        update_tag=update_tag,
+    load(
+        neo4j_session,
+        AzureSnapshotSchema(),
+        snapshots,
+        lastupdated=update_tag,
+        AZURE_SUBSCRIPTION_ID=subscription_id,
     )
 
 
 def cleanup_snapshot(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job(
-        "azure_import_snapshots_cleanup.json",
+    GraphJob.from_node_schema(AzureSnapshotSchema(), common_job_parameters).run(
         neo4j_session,
-        common_job_parameters,
     )
 
 
+def transform_vm_list(vm_list: List[Dict]) -> Tuple[List[Dict], List[Dict]]:
+    """
+    Transform the VM list to separate the VMs and their data disks.
+    """
+    transformed_vm_list = []
+    transformed_data_disk_list = []
+
+    for vm in vm_list:
+        for dd in vm.get("storage_profile", {}).get("data_disks", []):
+            dd["vm_id"] = vm["id"]
+            transformed_data_disk_list.append(dd)
+        transformed_vm_list.append(vm)
+
+    return transformed_vm_list, transformed_data_disk_list
+
+
 def sync_virtual_machine(
     neo4j_session: neo4j.Session,
     credentials: Credentials,
@@ -241,7 +181,14 @@ def sync_virtual_machine(
     common_job_parameters: Dict,
 ) -> None:
     vm_list = get_vm_list(credentials, subscription_id)
-    load_vms(neo4j_session, subscription_id, vm_list, update_tag)
+    transformed_vm_list, transformed_data_disk_list = transform_vm_list(vm_list)
+    load_vms(neo4j_session, subscription_id, transformed_vm_list, update_tag)
+    load_vm_data_disks(
+        neo4j_session,
+        subscription_id,
+        transformed_data_disk_list,
+        update_tag,
+    )
     cleanup_virtual_machine(neo4j_session, common_job_parameters)
 
 

cartography/intel/azure/container_instances.py

@@ -0,0 +1,95 @@
+import logging
+from typing import Any
+
+import neo4j
+from azure.core.exceptions import ClientAuthenticationError
+from azure.core.exceptions import HttpResponseError
+from azure.mgmt.containerinstance import ContainerInstanceManagementClient
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.azure.container_instance import AzureContainerInstanceSchema
+from cartography.util import timeit
+
+from .util.credentials import Credentials
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def get_container_instances(
+    credentials: Credentials, subscription_id: str
+) -> list[dict]:
+    try:
+        client = ContainerInstanceManagementClient(
+            credentials.credential, subscription_id
+        )
+        # NOTE: Azure Container Instances are called "Container Groups" in the SDK
+        return [cg.as_dict() for cg in client.container_groups.list()]
+    except (ClientAuthenticationError, HttpResponseError) as e:
+        logger.warning(
+            f"Failed to get Container Instances for subscription {subscription_id}: {str(e)}"
+        )
+        return []
+
+
+def transform_container_instances(container_groups: list[dict]) -> list[dict]:
+    transformed_instances: list[dict[str, Any]] = []
+    for group in container_groups:
+        transformed_instance = {
+            "id": group.get("id"),
+            "name": group.get("name"),
+            "location": group.get("location"),
+            "type": group.get("type"),
+            "provisioning_state": group.get("properties", {}).get("provisioning_state"),
+            "ip_address": ((group.get("properties") or {}).get("ip_address") or {}).get(
+                "ip"
+            ),
+            "os_type": group.get("properties", {}).get("os_type"),
+        }
+        transformed_instances.append(transformed_instance)
+    return transformed_instances
+
+
+@timeit
+def load_container_instances(
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    subscription_id: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        AzureContainerInstanceSchema(),
+        data,
+        lastupdated=update_tag,
+        AZURE_SUBSCRIPTION_ID=subscription_id,
+    )
+
+
+@timeit
+def cleanup_container_instances(
+    neo4j_session: neo4j.Session, common_job_parameters: dict
+) -> None:
+    GraphJob.from_node_schema(
+        AzureContainerInstanceSchema(), common_job_parameters
+    ).run(neo4j_session)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    credentials: Credentials,
+    subscription_id: str,
+    update_tag: int,
+    common_job_parameters: dict,
+) -> None:
+    logger.info(
+        f"Syncing Azure Container Instances for subscription {subscription_id}."
+    )
+    raw_groups = get_container_instances(credentials, subscription_id)
+    transformed_groups = transform_container_instances(raw_groups)
+    load_container_instances(
+        neo4j_session, transformed_groups, subscription_id, update_tag
+    )
+    cleanup_container_instances(neo4j_session, common_job_parameters)