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

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)

cartography/intel/aws/ec2/load_balancers.py

@@ -3,7 +3,6 @@ import logging
 import boto3
 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.load_balancer_listeners import ELBListenerSchema
@@ -11,6 +10,8 @@ from cartography.models.aws.ec2.load_balancers import LoadBalancerSchema
 from cartography.util import aws_handle_regions
 from cartography.util import timeit
 
+from .util import get_botocore_config
+
 logger = logging.getLogger(__name__)
 
 
@@ -29,7 +30,9 @@ def _get_listener_id(load_balancer_id: str, port: int, protocol: str) -> str:
     return f"{load_balancer_id}{port}{protocol}"
 
 
-def transform_load_balancer_listener_data(load_balancer_id: str, listener_data: list[dict]) -> list[dict]:
+def transform_load_balancer_listener_data(
+    load_balancer_id: str, listener_data: list[dict]
+) -> list[dict]:
     """
     Transform load balancer listener data into a format suitable for cartography ingestion.
 
@@ -42,21 +45,27 @@ def transform_load_balancer_listener_data(load_balancer_id: str, listener_data:
     """
     transformed = []
     for listener in listener_data:
-        listener_info = listener['Listener']
+        listener_info = listener["Listener"]
         transformed_listener = {
-            'id': _get_listener_id(load_balancer_id, listener_info['LoadBalancerPort'], listener_info['Protocol']),
-            'port': listener_info.get('LoadBalancerPort'),
-            'protocol': listener_info.get('Protocol'),
-            'instance_port': listener_info.get('InstancePort'),
-            'instance_protocol': listener_info.get('InstanceProtocol'),
-            'policy_names': listener.get('PolicyNames', []),
-            'LoadBalancerId': load_balancer_id,
+            "id": _get_listener_id(
+                load_balancer_id,
+                listener_info["LoadBalancerPort"],
+                listener_info["Protocol"],
+            ),
+            "port": listener_info.get("LoadBalancerPort"),
+            "protocol": listener_info.get("Protocol"),
+            "instance_port": listener_info.get("InstancePort"),
+            "instance_protocol": listener_info.get("InstanceProtocol"),
+            "policy_names": listener.get("PolicyNames", []),
+            "LoadBalancerId": load_balancer_id,
         }
         transformed.append(transformed_listener)
     return transformed
 
 
-def transform_load_balancer_data(load_balancers: list[dict]) -> tuple[list[dict], list[dict]]:
+def transform_load_balancer_data(
+    load_balancers: list[dict],
+) -> tuple[list[dict], list[dict]]:
     """
     Transform load balancer data into a format suitable for cartography ingestion.
 
@@ -70,35 +79,38 @@ def transform_load_balancer_data(load_balancers: list[dict]) -> tuple[list[dict]
     listener_data = []
 
     for lb in load_balancers:
-        load_balancer_id = lb['DNSName']
+        load_balancer_id = lb["DNSName"]
         transformed_lb = {
-            'id': load_balancer_id,
-            'name': lb['LoadBalancerName'],
-            'dnsname': lb['DNSName'],
-            'canonicalhostedzonename': lb.get('CanonicalHostedZoneName'),
-            'canonicalhostedzonenameid': lb.get('CanonicalHostedZoneNameID'),
-            'scheme': lb.get('Scheme'),
-            'createdtime': str(lb['CreatedTime']),
-            'GROUP_NAME': lb.get('SourceSecurityGroup', {}).get('GroupName'),
-            'GROUP_IDS': [str(group) for group in lb.get('SecurityGroups', [])],
-            'INSTANCE_IDS': [instance['InstanceId'] for instance in lb.get('Instances', [])],
-            'LISTENER_IDS': [
+            "id": load_balancer_id,
+            "name": lb["LoadBalancerName"],
+            "dnsname": lb["DNSName"],
+            "canonicalhostedzonename": lb.get("CanonicalHostedZoneName"),
+            "canonicalhostedzonenameid": lb.get("CanonicalHostedZoneNameID"),
+            "scheme": lb.get("Scheme"),
+            "createdtime": str(lb["CreatedTime"]),
+            "GROUP_NAME": lb.get("SourceSecurityGroup", {}).get("GroupName"),
+            "GROUP_IDS": [str(group) for group in lb.get("SecurityGroups", [])],
+            "INSTANCE_IDS": [
+                instance["InstanceId"] for instance in lb.get("Instances", [])
+            ],
+            "LISTENER_IDS": [
                 _get_listener_id(
                     load_balancer_id,
-                    listener['Listener']['LoadBalancerPort'],
-                    listener['Listener']['Protocol'],
-                ) for listener in lb.get('ListenerDescriptions', [])
+                    listener["Listener"]["LoadBalancerPort"],
+                    listener["Listener"]["Protocol"],
+                )
+                for listener in lb.get("ListenerDescriptions", [])
             ],
         }
         transformed.append(transformed_lb)
 
         # Classic ELB listeners are not returned anywhere else in AWS, so we must parse them out
         # of the describe_load_balancers response.
-        if lb.get('ListenerDescriptions'):
+        if lb.get("ListenerDescriptions"):
             listener_data.extend(
                 transform_load_balancer_listener_data(
                     load_balancer_id,
-                    lb.get('ListenerDescriptions', []),
+                    lb.get("ListenerDescriptions", []),
                 ),
             )
 
@@ -107,18 +119,25 @@ def transform_load_balancer_data(load_balancers: list[dict]) -> tuple[list[dict]
 
 @timeit
 @aws_handle_regions
-def get_loadbalancer_data(boto3_session: boto3.session.Session, region: str) -> list[dict]:
-    client = boto3_session.client('elb', region_name=region, config=get_botocore_config())
-    paginator = client.get_paginator('describe_load_balancers')
+def get_loadbalancer_data(
+    boto3_session: boto3.session.Session, region: str
+) -> list[dict]:
+    client = boto3_session.client(
+        "elb", region_name=region, config=get_botocore_config()
+    )
+    paginator = client.get_paginator("describe_load_balancers")
     elbs: list[dict] = []
     for page in paginator.paginate():
-        elbs.extend(page['LoadBalancerDescriptions'])
+        elbs.extend(page["LoadBalancerDescriptions"])
     return elbs
 
 
 @timeit
 def load_load_balancers(
-    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:
     load(
@@ -133,7 +152,10 @@ def load_load_balancers(
 
 @timeit
 def load_load_balancer_listeners(
-    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:
     load(
@@ -147,22 +169,40 @@ def load_load_balancer_listeners(
 
 
 @timeit
-def cleanup_load_balancers(neo4j_session: neo4j.Session, common_job_parameters: dict) -> None:
-    GraphJob.from_node_schema(ELBListenerSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(LoadBalancerSchema(), common_job_parameters).run(neo4j_session)
+def cleanup_load_balancers(
+    neo4j_session: neo4j.Session, common_job_parameters: dict
+) -> None:
+    GraphJob.from_node_schema(ELBListenerSchema(), common_job_parameters).run(
+        neo4j_session
+    )
+    GraphJob.from_node_schema(LoadBalancerSchema(), common_job_parameters).run(
+        neo4j_session
+    )
 
 
 @timeit
 def sync_load_balancers(
-    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 for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 load balancers for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_loadbalancer_data(boto3_session, region)
         transformed_data, listener_data = transform_load_balancer_data(data)
 
-        load_load_balancers(neo4j_session, transformed_data, region, current_aws_account_id, update_tag)
-        load_load_balancer_listeners(neo4j_session, listener_data, region, current_aws_account_id, update_tag)
+        load_load_balancers(
+            neo4j_session, transformed_data, region, current_aws_account_id, update_tag
+        )
+        load_load_balancer_listeners(
+            neo4j_session, listener_data, region, current_aws_account_id, update_tag
+        )
 
     cleanup_load_balancers(neo4j_session, common_job_parameters)