cartography 0.93.0rc1__py3-none-any.whl → 0.123.0__py3-none-any.whl

cartography/intel/azure/aks.py

@@ -0,0 +1,175 @@
+import logging
+from typing import Any
+
+import neo4j
+from azure.core.exceptions import ClientAuthenticationError
+from azure.core.exceptions import HttpResponseError
+from azure.mgmt.containerservice import ContainerServiceClient
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.azure.aks_cluster import AzureKubernetesClusterSchema
+from cartography.models.azure.aks_nodepool import AzureKubernetesNodePoolSchema
+from cartography.util import timeit
+
+from .util.credentials import Credentials
+
+logger = logging.getLogger(__name__)
+
+
+def _get_resource_group_from_id(resource_id: str) -> str:
+    """
+    Helper function to parse the resource group name from a full resource ID string.
+    """
+    # CORRECTED: Convert to lowercase to handle case inconsistencies from the API.
+    parts = resource_id.lower().split("/")
+    try:
+        # The resource group name is always the string after 'resourcegroups'
+        rg_index = parts.index("resourcegroups")
+        return parts[rg_index + 1]
+    except (ValueError, IndexError):
+        logger.warning(
+            f"Could not parse resource group name from resource ID: {resource_id}"
+        )
+        return ""
+
+
+@timeit
+def get_aks_clusters(
+    client: ContainerServiceClient, subscription_id: str
+) -> list[dict]:
+    try:
+        return [cluster.as_dict() for cluster in client.managed_clusters.list()]
+    except (ClientAuthenticationError, HttpResponseError) as e:
+        logger.warning(
+            f"Failed to get AKS clusters for subscription {subscription_id}: {str(e)}"
+        )
+        return []
+
+
+@timeit
+def get_agent_pools(
+    client: ContainerServiceClient,
+    cluster_name: str,
+    resource_group_name: str,
+) -> list[dict]:
+    try:
+        return [
+            pool.as_dict()
+            for pool in client.agent_pools.list(resource_group_name, cluster_name)
+        ]
+    except (ClientAuthenticationError, HttpResponseError) as e:
+        logger.warning(
+            f"Failed to get agent pools for cluster {cluster_name}: {str(e)}"
+        )
+        return []
+
+
+@timeit
+def transform_aks_clusters(clusters: list[dict]) -> list[dict]:
+    transformed_clusters: list[dict[str, Any]] = []
+    for cluster in clusters:
+        transformed_cluster = {
+            "id": cluster.get("id"),
+            "name": cluster.get("name"),
+            "location": cluster.get("location"),
+            "provisioning_state": cluster.get("properties", {}).get(
+                "provisioning_state"
+            ),
+            "kubernetes_version": cluster.get("properties", {}).get(
+                "kubernetes_version"
+            ),
+            "fqdn": cluster.get("properties", {}).get("fqdn"),
+        }
+        transformed_clusters.append(transformed_cluster)
+    return transformed_clusters
+
+
+@timeit
+def transform_agent_pools(agent_pools: list[dict]) -> list[dict]:
+    transformed_pools: list[dict[str, Any]] = []
+    for pool in agent_pools:
+        transformed_pool = {
+            "id": pool.get("id"),
+            "name": pool.get("name"),
+            "provisioning_state": pool.get("properties", {}).get("provisioning_state"),
+            "vm_size": pool.get("properties", {}).get("vm_size"),
+            "os_type": pool.get("properties", {}).get("os_type"),
+            "count": pool.get("properties", {}).get("count"),
+        }
+        transformed_pools.append(transformed_pool)
+    return transformed_pools
+
+
+@timeit
+def load_aks_clusters(
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    subscription_id: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        AzureKubernetesClusterSchema(),
+        data,
+        lastupdated=update_tag,
+        AZURE_SUBSCRIPTION_ID=subscription_id,
+    )
+
+
+@timeit
+def load_agent_pools(
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    cluster_id: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        AzureKubernetesNodePoolSchema(),
+        data,
+        lastupdated=update_tag,
+        CLUSTER_ID=cluster_id,
+    )
+
+
+@timeit
+def cleanup_clusters(neo4j_session: neo4j.Session, common_job_parameters: dict) -> None:
+    GraphJob.from_node_schema(
+        AzureKubernetesClusterSchema(), 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 Kubernetes Service for subscription {subscription_id}.")
+    client = ContainerServiceClient(credentials.credential, subscription_id)
+
+    clusters = get_aks_clusters(client, subscription_id)
+    transformed_clusters = transform_aks_clusters(clusters)
+    load_aks_clusters(neo4j_session, transformed_clusters, subscription_id, update_tag)
+
+    for cluster in clusters:
+        cluster_id = cluster.get("id")
+        if not cluster_id:
+            continue
+
+        resource_group_name = _get_resource_group_from_id(cluster_id)
+        if resource_group_name:
+            agent_pools = get_agent_pools(client, cluster["name"], resource_group_name)
+            transformed_pools = transform_agent_pools(agent_pools)
+            load_agent_pools(neo4j_session, transformed_pools, cluster_id, update_tag)
+
+            pool_cleanup_params = common_job_parameters.copy()
+            pool_cleanup_params["CLUSTER_ID"] = cluster_id
+            GraphJob.from_node_schema(
+                AzureKubernetesNodePoolSchema(), pool_cleanup_params
+            ).run(neo4j_session)
+
+    cleanup_clusters(neo4j_session, common_job_parameters)

cartography/intel/azure/app_service.py

@@ -0,0 +1,105 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+from azure.core.exceptions import ClientAuthenticationError
+from azure.core.exceptions import HttpResponseError
+from azure.mgmt.web import WebSiteManagementClient
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.models.azure.app_service import AzureAppServiceSchema
+from cartography.util import timeit
+
+from .util.credentials import Credentials
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def get_app_services(credentials: Credentials, subscription_id: str) -> List[Dict]:
+    """
+    Get a list of App Services from the given Azure subscription.
+    """
+    try:
+        client = WebSiteManagementClient(credentials.credential, subscription_id)
+        # NOTE: This is the same API call as Functions. We get all web apps
+        # and then filter them in the transform stage.
+        return [app.as_dict() for app in client.web_apps.list()]
+    except (ClientAuthenticationError, HttpResponseError) as e:
+        logger.warning(
+            f"Failed to get app services for subscription {subscription_id}: {str(e)}"
+        )
+        return []
+
+
+@timeit
+def transform_app_services(app_services_response: List[Dict]) -> List[Dict]:
+    """
+    Transform the raw API response to the dictionary structure that the model expects.
+    """
+    transformed_apps: List[Dict[str, Any]] = []
+    for app in app_services_response:
+        if "functionapp" not in app.get("kind", ""):
+            transformed_app = {
+                "id": app.get("id"),
+                "name": app.get("name"),
+                "kind": app.get("kind"),
+                "location": app.get("location"),
+                "state": app.get("state"),
+                "default_host_name": app.get("default_host_name"),
+                "https_only": app.get("https_only"),
+            }
+            transformed_apps.append(transformed_app)
+    return transformed_apps
+
+
+@timeit
+def load_app_services(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    subscription_id: str,
+    update_tag: int,
+) -> None:
+    """
+    Load the transformed Azure App Service data to Neo4j.
+    """
+    load(
+        neo4j_session,
+        AzureAppServiceSchema(),
+        data,
+        lastupdated=update_tag,
+        AZURE_SUBSCRIPTION_ID=subscription_id,
+    )
+
+
+@timeit
+def cleanup_app_services(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict
+) -> None:
+    """
+    Run the cleanup job for Azure App Services.
+    """
+    GraphJob.from_node_schema(AzureAppServiceSchema(), 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:
+    """
+    The main sync function for Azure App Services.
+    """
+    logger.info(f"Syncing Azure App Services for subscription {subscription_id}.")
+    raw_apps = get_app_services(credentials, subscription_id)
+    transformed_apps = transform_app_services(raw_apps)
+    load_app_services(neo4j_session, transformed_apps, subscription_id, update_tag)
+    cleanup_app_services(neo4j_session, common_job_parameters)

cartography/intel/azure/compute.py

@@ -1,19 +1,29 @@
 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 .util.credentials import Credentials
-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
+
 logger = logging.getLogger(__name__)
 
 
-def get_client(credentials: Credentials, subscription_id: str) -> ComputeManagementClient:
+def get_client(
+    credentials: Credentials,
+    subscription_id: str,
+) -> ComputeManagementClient:
     client = ComputeManagementClient(credentials, subscription_id)
     return client
 
@@ -24,8 +34,8 @@ def get_vm_list(credentials: Credentials, subscription_id: str) -> List[Dict]:
         vm_list = list(map(lambda x: x.as_dict(), client.virtual_machines.list_all()))
 
         for vm in vm_list:
-            x = vm['id'].split('/')
-            vm['resource_group'] = x[x.index('resourceGroups') + 1]
+            x = vm["id"].split("/")
+            vm["resource_group"] = x[x.index("resourceGroups") + 1]
 
         return vm_list
 
@@ -34,66 +44,43 @@ def get_vm_list(credentials: Credentials, subscription_id: str) -> List[Dict]:
         return []
 
 
-def load_vms(neo4j_session: neo4j.Session, subscription_id: str, 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,
+def load_vms(
+    neo4j_session: neo4j.Session,
+    subscription_id: str,
+    vm_list: List[Dict],
+    update_tag: int,
+) -> None:
+    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, 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,
+def load_vm_data_disks(
+    neo4j_session: neo4j.Session,
+    subscription_id: str,
+    data_disks: List[Dict],
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        AzureDataDiskSchema(),
+        data_disks,
+        lastupdated=update_tag,
+        AZURE_SUBSCRIPTION_ID=subscription_id,
     )
 
 
-def cleanup_virtual_machine(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('azure_import_virtual_machines_cleanup.json', neo4j_session, common_job_parameters)
+def cleanup_virtual_machine(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
+    GraphJob.from_node_schema(AzureVirtualMachineSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
 
 
 def get_disks(credentials: Credentials, subscription_id: str) -> List[Dict]:
@@ -102,8 +89,8 @@ def get_disks(credentials: Credentials, subscription_id: str) -> List[Dict]:
         disk_list = list(map(lambda x: x.as_dict(), client.disks.list()))
 
         for disk in disk_list:
-            x = disk['id'].split('/')
-            disk['resource_group'] = x[x.index('resourceGroups') + 1]
+            x = disk["id"].split("/")
+            disk["resource_group"] = x[x.index("resourceGroups") + 1]
 
         return disk_list
 
@@ -112,35 +99,25 @@ def get_disks(credentials: Credentials, subscription_id: str) -> List[Dict]:
         return []
 
 
-def load_disks(neo4j_session: neo4j.Session, subscription_id: str, 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,
+def load_disks(
+    neo4j_session: neo4j.Session,
+    subscription_id: str,
+    disk_list: List[Dict],
+    update_tag: int,
+) -> None:
+    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', neo4j_session, common_job_parameters)
+    GraphJob.from_node_schema(AzureDiskSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
 
 
 def get_snapshots_list(credentials: Credentials, subscription_id: str) -> List[Dict]:
@@ -149,8 +126,8 @@ def get_snapshots_list(credentials: Credentials, subscription_id: str) -> List[D
         snapshots = list(map(lambda x: x.as_dict(), client.snapshots.list()))
 
         for snapshot in snapshots:
-            x = snapshot['id'].split('/')
-            snapshot['resource_group'] = x[x.index('resourceGroups') + 1]
+            x = snapshot["id"].split("/")
+            snapshot["resource_group"] = x[x.index("resourceGroups") + 1]
 
         return snapshots
 
@@ -159,47 +136,67 @@ def get_snapshots_list(credentials: Credentials, subscription_id: str) -> List[D
         return []
 
 
-def load_snapshots(neo4j_session: neo4j.Session, subscription_id: str, 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,
+def load_snapshots(
+    neo4j_session: neo4j.Session,
+    subscription_id: str,
+    snapshots: List[Dict],
+    update_tag: int,
+) -> None:
+    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', neo4j_session, common_job_parameters)
+    GraphJob.from_node_schema(AzureSnapshotSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+
+
+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, subscription_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    credentials: Credentials,
+    subscription_id: str,
+    update_tag: int,
     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)
 
 
 def sync_disk(
-    neo4j_session: neo4j.Session, credentials: Credentials, subscription_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    credentials: Credentials,
+    subscription_id: str,
+    update_tag: int,
     common_job_parameters: Dict,
 ) -> None:
     disk_list = get_disks(credentials, subscription_id)
@@ -208,7 +205,10 @@ def sync_disk(
 
 
 def sync_snapshot(
-    neo4j_session: neo4j.Session, credentials: Credentials, subscription_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    credentials: Credentials,
+    subscription_id: str,
+    update_tag: int,
     common_job_parameters: Dict,
 ) -> None:
     snapshots = get_snapshots_list(credentials, subscription_id)
@@ -218,11 +218,32 @@ def sync_snapshot(
 
 @timeit
 def sync(
-    neo4j_session: neo4j.Session, credentials: Credentials, subscription_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    credentials: Credentials,
+    subscription_id: str,
+    update_tag: int,
     common_job_parameters: Dict,
 ) -> None:
     logger.info("Syncing VM for subscription '%s'.", subscription_id)
 
-    sync_virtual_machine(neo4j_session, credentials, subscription_id, update_tag, common_job_parameters)
-    sync_disk(neo4j_session, credentials, subscription_id, update_tag, common_job_parameters)
-    sync_snapshot(neo4j_session, credentials, subscription_id, update_tag, common_job_parameters)
+    sync_virtual_machine(
+        neo4j_session,
+        credentials,
+        subscription_id,
+        update_tag,
+        common_job_parameters,
+    )
+    sync_disk(
+        neo4j_session,
+        credentials,
+        subscription_id,
+        update_tag,
+        common_job_parameters,
+    )
+    sync_snapshot(
+        neo4j_session,
+        credentials,
+        subscription_id,
+        update_tag,
+        common_job_parameters,
+    )