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

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)

cartography/intel/aws/ec2/load_balancer_v2s.py

@@ -6,38 +6,47 @@ import boto3
 import botocore
 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_load_balancer_v2_listeners(client: botocore.client.BaseClient, load_balancer_arn: str) -> List[Dict]:
-    paginator = client.get_paginator('describe_listeners')
+def get_load_balancer_v2_listeners(
+    client: botocore.client.BaseClient,
+    load_balancer_arn: str,
+) -> List[Dict]:
+    paginator = client.get_paginator("describe_listeners")
     listeners: List[Dict] = []
     for page in paginator.paginate(LoadBalancerArn=load_balancer_arn):
-        listeners.extend(page['Listeners'])
+        listeners.extend(page["Listeners"])
 
     return listeners
 
 
 @timeit
-def get_load_balancer_v2_target_groups(client: botocore.client.BaseClient, load_balancer_arn: str) -> List[Dict]:
-    paginator = client.get_paginator('describe_target_groups')
+def get_load_balancer_v2_target_groups(
+    client: botocore.client.BaseClient,
+    load_balancer_arn: str,
+) -> List[Dict]:
+    paginator = client.get_paginator("describe_target_groups")
     target_groups: List[Dict] = []
     for page in paginator.paginate(LoadBalancerArn=load_balancer_arn):
-        target_groups.extend(page['TargetGroups'])
+        target_groups.extend(page["TargetGroups"])
 
     # Add instance data
     for target_group in target_groups:
-        target_group['Targets'] = []
-        target_health = client.describe_target_health(TargetGroupArn=target_group['TargetGroupArn'])
-        for target_health_description in target_health['TargetHealthDescriptions']:
-            target_group['Targets'].append(target_health_description['Target']['Id'])
+        target_group["Targets"] = []
+        target_health = client.describe_target_health(
+            TargetGroupArn=target_group["TargetGroupArn"],
+        )
+        for target_health_description in target_health["TargetHealthDescriptions"]:
+            target_group["Targets"].append(target_health_description["Target"]["Id"])
 
     return target_groups
 
@@ -45,22 +54,35 @@ def get_load_balancer_v2_target_groups(client: botocore.client.BaseClient, load_
 @timeit
 @aws_handle_regions
 def get_loadbalancer_v2_data(boto3_session: boto3.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('elbv2', region_name=region, config=get_botocore_config())
-    paginator = client.get_paginator('describe_load_balancers')
+    client = boto3_session.client(
+        "elbv2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    paginator = client.get_paginator("describe_load_balancers")
     elbv2s: List[Dict] = []
     for page in paginator.paginate():
-        elbv2s.extend(page['LoadBalancers'])
+        elbv2s.extend(page["LoadBalancers"])
 
     # Make extra calls to get listeners
     for elbv2 in elbv2s:
-        elbv2['Listeners'] = get_load_balancer_v2_listeners(client, elbv2['LoadBalancerArn'])
-        elbv2['TargetGroups'] = get_load_balancer_v2_target_groups(client, elbv2['LoadBalancerArn'])
+        elbv2["Listeners"] = get_load_balancer_v2_listeners(
+            client,
+            elbv2["LoadBalancerArn"],
+        )
+        elbv2["TargetGroups"] = get_load_balancer_v2_target_groups(
+            client,
+            elbv2["LoadBalancerArn"],
+        )
     return elbv2s
 
 
 @timeit
 def load_load_balancer_v2s(
-    neo4j_session: neo4j.Session, data: List[Dict], region: str, current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    current_aws_account_id: str,
     update_tag: int,
 ) -> None:
     ingest_load_balancer_v2 = """
@@ -95,10 +117,16 @@ def load_load_balancer_v2s(
 
         if lb["AvailabilityZones"]:
             az = lb["AvailabilityZones"]
-            load_load_balancer_v2_subnets(neo4j_session, load_balancer_id, az, region, update_tag)
+            load_load_balancer_v2_subnets(
+                neo4j_session,
+                load_balancer_id,
+                az,
+                region,
+                update_tag,
+            )
 
         # NLB's don't have SecurityGroups, so check for one first.
-        if 'SecurityGroups' in lb and lb["SecurityGroups"]:
+        if "SecurityGroups" in lb and lb["SecurityGroups"]:
             ingest_load_balancer_v2_security_group = """
             MATCH (elbv2:LoadBalancerV2{id: $ID}),
             (group:EC2SecurityGroup{groupid: $GROUP_ID})
@@ -114,20 +142,31 @@ def load_load_balancer_v2s(
                     update_tag=update_tag,
                 )
 
-        if lb['Listeners']:
-            load_load_balancer_v2_listeners(neo4j_session, load_balancer_id, lb['Listeners'], update_tag)
+        if lb["Listeners"]:
+            load_load_balancer_v2_listeners(
+                neo4j_session,
+                load_balancer_id,
+                lb["Listeners"],
+                update_tag,
+            )
 
-        if lb['TargetGroups']:
+        if lb["TargetGroups"]:
             load_load_balancer_v2_target_groups(
-                neo4j_session, load_balancer_id, lb['TargetGroups'],
-                current_aws_account_id, update_tag,
+                neo4j_session,
+                load_balancer_id,
+                lb["TargetGroups"],
+                current_aws_account_id,
+                update_tag,
             )
 
 
 @timeit
 def load_load_balancer_v2_subnets(
-    neo4j_session: neo4j.Session, load_balancer_id: str, az_data: List[Dict],
-    region: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    load_balancer_id: str,
+    az_data: List[Dict],
+    region: str,
+    update_tag: int,
 ) -> None:
     ingest_load_balancer_subnet = """
     MATCH (elbv2:LoadBalancerV2{id: $ID})
@@ -143,7 +182,7 @@ def load_load_balancer_v2_subnets(
         neo4j_session.run(
             ingest_load_balancer_subnet,
             ID=load_balancer_id,
-            SubnetId=az['SubnetId'],
+            SubnetId=az["SubnetId"],
             region=region,
             update_tag=update_tag,
         )
@@ -151,7 +190,10 @@ def load_load_balancer_v2_subnets(
 
 @timeit
 def load_load_balancer_v2_target_groups(
-    neo4j_session: neo4j.Session, load_balancer_id: str, target_groups: List[Dict], current_aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    load_balancer_id: str,
+    target_groups: List[Dict],
+    current_aws_account_id: str,
     update_tag: int,
 ) -> None:
     ingest_instances = """
@@ -169,7 +211,7 @@ def load_load_balancer_v2_target_groups(
     """
     for target_group in target_groups:
 
-        if not target_group['TargetType'] == 'instance':
+        if not target_group["TargetType"] == "instance":
             # Only working on EC2 Instances now. TODO: Add IP & Lambda EXPOSE.
             continue
 
@@ -179,16 +221,18 @@ def load_load_balancer_v2_target_groups(
                 ID=load_balancer_id,
                 INSTANCE_ID=instance,
                 AWS_ACCOUNT_ID=current_aws_account_id,
-                TARGET_GROUP_ARN=target_group.get('TargetGroupArn'),
-                PORT=target_group.get('Port'),
-                PROTOCOL=target_group.get('Protocol'),
+                TARGET_GROUP_ARN=target_group.get("TargetGroupArn"),
+                PORT=target_group.get("Port"),
+                PROTOCOL=target_group.get("Protocol"),
                 update_tag=update_tag,
             )
 
 
 @timeit
 def load_load_balancer_v2_listeners(
-    neo4j_session: neo4j.Session, load_balancer_id: str, listener_data: List[Dict],
+    neo4j_session: neo4j.Session,
+    load_balancer_id: str,
+    listener_data: List[Dict],
     update_tag: int,
 ) -> None:
     ingest_listener = """
@@ -215,18 +259,39 @@ def load_load_balancer_v2_listeners(
 
 
 @timeit
-def cleanup_load_balancer_v2s(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
+def cleanup_load_balancer_v2s(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
     """Delete elbv2's and dependent resources in the DB without the most recent lastupdated tag."""
-    run_cleanup_job('aws_ingest_load_balancers_v2_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job(
+        "aws_ingest_load_balancers_v2_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
 
 
 @timeit
 def sync_load_balancer_v2s(
-    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 EC2 load balancers v2 for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 load balancers v2 for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_loadbalancer_v2_data(boto3_session, region)
-        load_load_balancer_v2s(neo4j_session, data, region, current_aws_account_id, update_tag)
+        load_load_balancer_v2s(
+            neo4j_session,
+            data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_load_balancer_v2s(neo4j_session, common_job_parameters)