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

cartography/intel/aws/ec2/internet_gateways.py

@@ -5,25 +5,36 @@ from typing import List
 import boto3
 import neo4j
 
-from .util import get_botocore_config
 from cartography.util import aws_handle_regions
 from cartography.util import run_cleanup_job
 from cartography.util import timeit
 
+from .util import get_botocore_config
+
 logger = logging.getLogger(__name__)
 
 
 @timeit
 @aws_handle_regions
-def get_internet_gateways(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    return client.describe_internet_gateways()['InternetGateways']
+def get_internet_gateways(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    return client.describe_internet_gateways()["InternetGateways"]
 
 
 @timeit
 def load_internet_gateways(
-    neo4j_session: neo4j.Session, internet_gateways: List[Dict], region: str,
-    current_aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    internet_gateways: List[Dict],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     logger.info("Loading %d Internet Gateways in %s.", len(internet_gateways), region)
     # TODO: Right now this won't work in non-AWS commercial (GovCloud, China) as partition is hardcoded
@@ -64,17 +75,35 @@ def load_internet_gateways(
 @timeit
 def cleanup(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
     logger.debug("Running Internet Gateway cleanup job.")
-    run_cleanup_job('aws_import_internet_gateways_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job(
+        "aws_import_internet_gateways_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 @timeit
 def sync_internet_gateways(
-    neo4j_session: neo4j.Session, boto3_session: boto3.session.Session, regions: List[str], current_aws_account_id: str,
-    update_tag: int, common_job_parameters: Dict,
+    neo4j_session: neo4j.Session,
+    boto3_session: boto3.session.Session,
+    regions: List[str],
+    current_aws_account_id: str,
+    update_tag: int,
+    common_job_parameters: Dict,
 ) -> None:
     for region in regions:
-        logger.info("Syncing Internet Gateways for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing Internet Gateways for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         internet_gateways = get_internet_gateways(boto3_session, region)
-        load_internet_gateways(neo4j_session, internet_gateways, region, current_aws_account_id, update_tag)
+        load_internet_gateways(
+            neo4j_session,
+            internet_gateways,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
 
     cleanup(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/key_pairs.py

@@ -17,34 +17,43 @@ logger = logging.getLogger(__name__)
 
 @timeit
 @aws_handle_regions
-def get_ec2_key_pairs(boto3_session: boto3.session.Session, region: str) -> list[dict[str, Any]]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    return client.describe_key_pairs()['KeyPairs']
+def get_ec2_key_pairs(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> list[dict[str, Any]]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    return client.describe_key_pairs()["KeyPairs"]
 
 
 def transform_ec2_key_pairs(
-        key_pairs: list[dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
+    key_pairs: list[dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
 ) -> list[dict[str, Any]]:
     transformed_key_pairs = []
     for key_pair in key_pairs:
         key_name = key_pair["KeyName"]
-        transformed_key_pairs.append({
-            'KeyPairArn': f'arn:aws:ec2:{region}:{current_aws_account_id}:key-pair/{key_name}',
-            'KeyName': key_name,
-            'KeyFingerprint': key_pair.get("KeyFingerprint"),
-        })
+        transformed_key_pairs.append(
+            {
+                "KeyPairArn": f"arn:aws:ec2:{region}:{current_aws_account_id}:key-pair/{key_name}",
+                "KeyName": key_name,
+                "KeyFingerprint": key_pair.get("KeyFingerprint"),
+            },
+        )
     return transformed_key_pairs
 
 
 @timeit
 def load_ec2_key_pairs(
-        neo4j_session: neo4j.Session,
-        data: list[dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     # Load EC2 keypairs as known by describe-key-pairs
     logger.info(f"Loading {len(data)} EC2 keypairs for region '{region}' into graph.")
@@ -59,8 +68,13 @@ def load_ec2_key_pairs(
 
 
 @timeit
-def cleanup_ec2_key_pairs(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    GraphJob.from_node_schema(EC2KeyPairSchema(), common_job_parameters).run(neo4j_session)
+def cleanup_ec2_key_pairs(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
+    GraphJob.from_node_schema(EC2KeyPairSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
 
 
 @timeit
@@ -73,8 +87,18 @@ def sync_ec2_key_pairs(
     common_job_parameters: dict[str, Any],
 ) -> None:
     for region in regions:
-        logger.info("Syncing EC2 key pairs for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 key pairs for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_ec2_key_pairs(boto3_session, region)
         transformed_data = transform_ec2_key_pairs(data, region, current_aws_account_id)
-        load_ec2_key_pairs(neo4j_session, transformed_data, region, current_aws_account_id, update_tag)
+        load_ec2_key_pairs(
+            neo4j_session,
+            transformed_data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_ec2_key_pairs(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/launch_templates.py

@@ -5,14 +5,17 @@ import boto3
 import botocore
 import neo4j
 
-from .util import get_botocore_config
 from cartography.client.core.tx import load
 from cartography.graph.job import GraphJob
-from cartography.models.aws.ec2.launch_template_versions import LaunchTemplateVersionSchema
+from cartography.models.aws.ec2.launch_template_versions import (
+    LaunchTemplateVersionSchema,
+)
 from cartography.models.aws.ec2.launch_templates import LaunchTemplateSchema
 from cartography.util import aws_handle_regions
 from cartography.util import timeit
 
+from .util import get_botocore_config
+
 logger = logging.getLogger(__name__)
 
 
@@ -22,11 +25,13 @@ def get_launch_templates(
     boto3_session: boto3.session.Session,
     region: str,
 ) -> list[dict[str, Any]]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    paginator = client.get_paginator('describe_launch_templates')
+    client = boto3_session.client(
+        "ec2", region_name=region, config=get_botocore_config()
+    )
+    paginator = client.get_paginator("describe_launch_templates")
     templates: list[dict[str, Any]] = []
     for page in paginator.paginate():
-        paginated_templates = page['LaunchTemplates']
+        paginated_templates = page["LaunchTemplates"]
         templates.extend(paginated_templates)
     return templates
 
@@ -40,8 +45,10 @@ def get_launch_template_versions(
 ) -> list[dict[str, Any]]:
     template_versions: list[dict[str, Any]] = []
     for template in launch_templates:
-        launch_template_id = template['LaunchTemplateId']
-        versions = get_launch_template_versions_by_template(boto3_session, launch_template_id, region)
+        launch_template_id = template["LaunchTemplateId"]
+        versions = get_launch_template_versions_by_template(
+            boto3_session, launch_template_id, region
+        )
         template_versions.extend(versions)
 
     return template_versions
@@ -50,46 +57,54 @@ def get_launch_template_versions(
 @timeit
 @aws_handle_regions
 def get_launch_template_versions_by_template(
-        boto3_session: boto3.session.Session,
-        launch_template_id: str,
-        region: str,
+    boto3_session: boto3.session.Session,
+    launch_template_id: str,
+    region: str,
 ) -> list[dict[str, Any]]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    v_paginator = client.get_paginator('describe_launch_template_versions')
+    client = boto3_session.client(
+        "ec2", region_name=region, config=get_botocore_config()
+    )
+    v_paginator = client.get_paginator("describe_launch_template_versions")
     template_versions = []
     try:
         for versions in v_paginator.paginate(LaunchTemplateId=launch_template_id):
-            template_versions.extend(versions['LaunchTemplateVersions'])
+            template_versions.extend(versions["LaunchTemplateVersions"])
     except botocore.exceptions.ClientError as e:
-        error_code = e.response['Error']['Code']
-        if error_code == 'InvalidLaunchTemplateId.NotFound':
-            logger.warning("Launch template %s no longer exists in region %s", launch_template_id, region)
+        error_code = e.response["Error"]["Code"]
+        if error_code == "InvalidLaunchTemplateId.NotFound":
+            logger.warning(
+                "Launch template %s no longer exists in region %s",
+                launch_template_id,
+                region,
+            )
         else:
             raise
     return template_versions
 
 
-def transform_launch_templates(templates: list[dict[str, Any]], versions: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    valid_template_ids = {v['LaunchTemplateId'] for v in versions}
+def transform_launch_templates(
+    templates: list[dict[str, Any]], versions: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
+    valid_template_ids = {v["LaunchTemplateId"] for v in versions}
     result: list[dict[str, Any]] = []
     for template in templates:
-        if template['LaunchTemplateId'] not in valid_template_ids:
+        if template["LaunchTemplateId"] not in valid_template_ids:
             continue
 
         current = template.copy()
         # Convert CreateTime to timestamp string
-        current['CreateTime'] = str(int(current['CreateTime'].timestamp()))
+        current["CreateTime"] = str(int(current["CreateTime"].timestamp()))
         result.append(current)
     return result
 
 
 @timeit
 def load_launch_templates(
-        neo4j_session: neo4j.Session,
-        data: list[dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     load(
         neo4j_session,
@@ -101,41 +116,45 @@ def load_launch_templates(
     )
 
 
-def transform_launch_template_versions(versions: list[dict[str, Any]]) -> list[dict[str, Any]]:
+def transform_launch_template_versions(
+    versions: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
     result: list[dict[str, Any]] = []
     for version in versions:
         current = version.copy()
 
         # Reformat some fields
-        current['Id'] = f"{version['LaunchTemplateId']}-{version['VersionNumber']}"
-        current['CreateTime'] = str(int(version['CreateTime'].timestamp()))
+        current["Id"] = f"{version['LaunchTemplateId']}-{version['VersionNumber']}"
+        current["CreateTime"] = str(int(version["CreateTime"].timestamp()))
 
         # Handle the nested object returned from boto
-        ltd = version['LaunchTemplateData']
-        current['KernelId'] = ltd.get('KernelId')
-        current['EbsOptimized'] = ltd.get('EbsOptimized')
-        current['IamInstanceProfileArn'] = ltd.get('IamInstanceProfileArn')
-        current['IamInstanceProfileName'] = ltd.get('IamInstanceProfileName')
-        current['ImageId'] = ltd.get('ImageId')
-        current['InstanceType'] = ltd.get('InstanceType')
-        current['KeyName'] = ltd.get('KeyName')
-        current['MonitoringEnabled'] = ltd.get('MonitoringEnabled')
-        current['RamdiskId'] = ltd.get('RamdiskId')
-        current['DisableApiTermination'] = ltd.get('DisableApiTermination')
-        current['InstanceInitiatedShutDownBehavior'] = ltd.get('InstanceInitiatedShutDownBehavior')
-        current['SecurityGroupIds'] = ltd.get('SecurityGroupIds')
-        current['SecurityGroups'] = ltd.get('SecurityGroups')
+        ltd = version["LaunchTemplateData"]
+        current["KernelId"] = ltd.get("KernelId")
+        current["EbsOptimized"] = ltd.get("EbsOptimized")
+        current["IamInstanceProfileArn"] = ltd.get("IamInstanceProfileArn")
+        current["IamInstanceProfileName"] = ltd.get("IamInstanceProfileName")
+        current["ImageId"] = ltd.get("ImageId")
+        current["InstanceType"] = ltd.get("InstanceType")
+        current["KeyName"] = ltd.get("KeyName")
+        current["MonitoringEnabled"] = ltd.get("MonitoringEnabled")
+        current["RamdiskId"] = ltd.get("RamdiskId")
+        current["DisableApiTermination"] = ltd.get("DisableApiTermination")
+        current["InstanceInitiatedShutDownBehavior"] = ltd.get(
+            "InstanceInitiatedShutDownBehavior",
+        )
+        current["SecurityGroupIds"] = ltd.get("SecurityGroupIds")
+        current["SecurityGroups"] = ltd.get("SecurityGroups")
         result.append(current)
     return result
 
 
 @timeit
 def load_launch_template_versions(
-        neo4j_session: neo4j.Session,
-        data: list[dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     load(
         neo4j_session,
@@ -148,35 +167,58 @@ def load_launch_template_versions(
 
 
 @timeit
-def cleanup(neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]) -> None:
+def cleanup(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
     logger.info("Running launch template cleanup job.")
-    cleanup_job = GraphJob.from_node_schema(LaunchTemplateSchema(), common_job_parameters)
+    cleanup_job = GraphJob.from_node_schema(
+        LaunchTemplateSchema(),
+        common_job_parameters,
+    )
     cleanup_job.run(neo4j_session)
 
-    cleanup_job = GraphJob.from_node_schema(LaunchTemplateVersionSchema(), common_job_parameters)
+    cleanup_job = GraphJob.from_node_schema(
+        LaunchTemplateVersionSchema(),
+        common_job_parameters,
+    )
     cleanup_job.run(neo4j_session)
 
 
 @timeit
 def sync_ec2_launch_templates(
-        neo4j_session: neo4j.Session,
-        boto3_session: boto3.session.Session,
-        regions: list[str],
-        current_aws_account_id: str,
-        update_tag: int,
-        common_job_parameters: dict[str, Any],
+    neo4j_session: neo4j.Session,
+    boto3_session: boto3.session.Session,
+    regions: list[str],
+    current_aws_account_id: str,
+    update_tag: int,
+    common_job_parameters: dict[str, Any],
 ) -> None:
     for region in regions:
-        logger.info(f"Syncing launch templates for region '{region}' in account '{current_aws_account_id}'.")
+        logger.info(
+            f"Syncing launch templates for region '{region}' in account '{current_aws_account_id}'."
+        )
         templates = get_launch_templates(boto3_session, region)
         versions = get_launch_template_versions(boto3_session, region, templates)
 
         # Transform and load the templates that have versions
         transformed_templates = transform_launch_templates(templates, versions)
-        load_launch_templates(neo4j_session, transformed_templates, region, current_aws_account_id, update_tag)
+        load_launch_templates(
+            neo4j_session,
+            transformed_templates,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
 
         # Transform and load the versions
         transformed_versions = transform_launch_template_versions(versions)
-        load_launch_template_versions(neo4j_session, transformed_versions, region, current_aws_account_id, update_tag)
+        load_launch_template_versions(
+            neo4j_session,
+            transformed_versions,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
 
     cleanup(neo4j_session, common_job_parameters)