cartography 0.102.0rc1__py3-none-any.whl → 0.103.0rc1__py3-none-any.whl

cartography/intel/azure/__init__.py

@@ -5,6 +5,9 @@ from typing import Optional
 
 import neo4j
 
+from cartography.config import Config
+from cartography.util import timeit
+
 from . import compute
 from . import cosmosdb
 from . import sql
@@ -13,24 +16,52 @@ from . import subscription
 from . import tenant
 from .util.credentials import Authenticator
 from .util.credentials import Credentials
-from cartography.config import Config
-from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
 
 
 def _sync_one_subscription(
-    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:
-    compute.sync(neo4j_session, credentials.arm_credentials, subscription_id, update_tag, common_job_parameters)
-    cosmosdb.sync(neo4j_session, credentials.arm_credentials, subscription_id, update_tag, common_job_parameters)
-    sql.sync(neo4j_session, credentials.arm_credentials, subscription_id, update_tag, common_job_parameters)
-    storage.sync(neo4j_session, credentials.arm_credentials, subscription_id, update_tag, common_job_parameters)
+    compute.sync(
+        neo4j_session,
+        credentials.arm_credentials,
+        subscription_id,
+        update_tag,
+        common_job_parameters,
+    )
+    cosmosdb.sync(
+        neo4j_session,
+        credentials.arm_credentials,
+        subscription_id,
+        update_tag,
+        common_job_parameters,
+    )
+    sql.sync(
+        neo4j_session,
+        credentials.arm_credentials,
+        subscription_id,
+        update_tag,
+        common_job_parameters,
+    )
+    storage.sync(
+        neo4j_session,
+        credentials.arm_credentials,
+        subscription_id,
+        update_tag,
+        common_job_parameters,
+    )
 
 
 def _sync_tenant(
-    neo4j_session: neo4j.Session, tenant_id: str, current_user: Optional[str], update_tag: int,
+    neo4j_session: neo4j.Session,
+    tenant_id: str,
+    current_user: Optional[str],
+    update_tag: int,
     common_job_parameters: Dict,
 ) -> None:
     logger.info("Syncing Azure Tenant: %s", tenant_id)
@@ -38,18 +69,34 @@ def _sync_tenant(
 
 
 def _sync_multiple_subscriptions(
-    neo4j_session: neo4j.Session, credentials: Credentials, tenant_id: str, subscriptions: List[Dict],
-    update_tag: int, common_job_parameters: Dict,
+    neo4j_session: neo4j.Session,
+    credentials: Credentials,
+    tenant_id: str,
+    subscriptions: List[Dict],
+    update_tag: int,
+    common_job_parameters: Dict,
 ) -> None:
     logger.info("Syncing Azure subscriptions")
 
-    subscription.sync(neo4j_session, tenant_id, subscriptions, update_tag, common_job_parameters)
+    subscription.sync(
+        neo4j_session,
+        tenant_id,
+        subscriptions,
+        update_tag,
+        common_job_parameters,
+    )
 
     for sub in subscriptions:
-        logger.info("Syncing Azure Subscription with ID '%s'", sub['subscriptionId'])
-        common_job_parameters['AZURE_SUBSCRIPTION_ID'] = sub['subscriptionId']
-
-        _sync_one_subscription(neo4j_session, credentials, sub['subscriptionId'], update_tag, common_job_parameters)
+        logger.info("Syncing Azure Subscription with ID '%s'", sub["subscriptionId"])
+        common_job_parameters["AZURE_SUBSCRIPTION_ID"] = sub["subscriptionId"]
+
+        _sync_one_subscription(
+            neo4j_session,
+            credentials,
+            sub["subscriptionId"],
+            update_tag,
+            common_job_parameters,
+        )
 
     del common_job_parameters["AZURE_SUBSCRIPTION_ID"]
 
@@ -64,7 +111,9 @@ def start_azure_ingestion(neo4j_session: neo4j.Session, config: Config) -> None:
     try:
         if config.azure_sp_auth:
             credentials = Authenticator().authenticate_sp(
-                config.azure_tenant_id, config.azure_client_id, config.azure_client_secret,
+                config.azure_tenant_id,
+                config.azure_client_id,
+                config.azure_client_secret,
             )
         else:
             credentials = Authenticator().authenticate_cli()
@@ -80,7 +129,10 @@ def start_azure_ingestion(neo4j_session: neo4j.Session, config: Config) -> None:
         return
 
     _sync_tenant(
-        neo4j_session, credentials.get_tenant_id(), credentials.get_current_user(), config.update_tag,
+        neo4j_session,
+        credentials.get_tenant_id(),
+        credentials.get_current_user(),
+        config.update_tag,
         common_job_parameters,
     )
 
@@ -88,7 +140,10 @@ def start_azure_ingestion(neo4j_session: neo4j.Session, config: Config) -> None:
         subscriptions = subscription.get_all_azure_subscriptions(credentials)
 
     else:
-        subscriptions = subscription.get_current_azure_subscription(credentials, credentials.subscription_id)
+        subscriptions = subscription.get_current_azure_subscription(
+            credentials,
+            credentials.subscription_id,
+        )
 
     if not subscriptions:
         logger.warning(
@@ -97,6 +152,10 @@ def start_azure_ingestion(neo4j_session: neo4j.Session, config: Config) -> None:
         return
 
     _sync_multiple_subscriptions(
-        neo4j_session, credentials, credentials.get_tenant_id(), subscriptions, config.update_tag,
+        neo4j_session,
+        credentials,
+        credentials.get_tenant_id(),
+        subscriptions,
+        config.update_tag,
         common_job_parameters,
     )

cartography/intel/azure/compute.py

@@ -6,14 +6,18 @@ 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.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 +28,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,7 +38,12 @@ 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:
+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})
@@ -62,11 +71,21 @@ def load_vms(neo4j_session: neo4j.Session, subscription_id: str, vm_list: List[D
     )
 
     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:
+        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})
@@ -92,8 +111,15 @@ def load_vm_data_disks(neo4j_session: neo4j.Session, vm_id: str, data_disks: Lis
     )
 
 
-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:
+    run_cleanup_job(
+        "azure_import_virtual_machines_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 def get_disks(credentials: Credentials, subscription_id: str) -> List[Dict]:
@@ -102,8 +128,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,7 +138,12 @@ 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:
+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})
@@ -140,7 +171,11 @@ def load_disks(neo4j_session: neo4j.Session, subscription_id: str, disk_list: Li
 
 
 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)
+    run_cleanup_job(
+        "azure_import_disks_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 def get_snapshots_list(credentials: Credentials, subscription_id: str) -> List[Dict]:
@@ -149,8 +184,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,7 +194,12 @@ 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:
+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})
@@ -186,11 +226,18 @@ def load_snapshots(neo4j_session: neo4j.Session, subscription_id: str, snapshots
 
 
 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)
+    run_cleanup_job(
+        "azure_import_snapshots_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 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)
@@ -199,7 +246,10 @@ def sync_virtual_machine(
 
 
 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 +258,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 +271,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,
+    )