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

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)

cartography/intel/aws/ec2/network_acls.py

@@ -5,7 +5,6 @@ from typing import Any
 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.network_acl_rules import EC2NetworkAclEgressRuleSchema
@@ -14,51 +13,63 @@ from cartography.models.aws.ec2.network_acls import EC2NetworkAclSchema
 from cartography.util import aws_handle_regions
 from cartography.util import timeit
 
+from .util import get_botocore_config
+
 logger = logging.getLogger(__name__)
 
 Ec2AclObjects = namedtuple(
-    "Ec2AclObjects", [
-        'network_acls',
-        'inbound_rules',
-        'outbound_rules',
+    "Ec2AclObjects",
+    [
+        "network_acls",
+        "inbound_rules",
+        "outbound_rules",
     ],
 )
 
 
 @timeit
 @aws_handle_regions
-def get_network_acl_data(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_network_acls')
+def get_network_acl_data(
+    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_network_acls")
     acls = []
     for page in paginator.paginate():
-        acls.extend(page['NetworkAcls'])
+        acls.extend(page["NetworkAcls"])
     return acls
 
 
 def transform_network_acl_data(
-        data_list: list[dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
+    data_list: list[dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
 ) -> Ec2AclObjects:
     network_acls = []
     inbound_rules = []
     outbound_rules = []
 
     for network_acl in data_list:
-        network_acl_id = network_acl['NetworkAclId']
+        network_acl_id = network_acl["NetworkAclId"]
         base_network_acl = {
-            'Id': network_acl_id,
-            'Arn': f'arn:aws:ec2:{region}:{current_aws_account_id}:network-acl/{network_acl_id}',
-            'IsDefault': network_acl['IsDefault'],
-            'VpcId': network_acl['VpcId'],
-            'OwnerId': network_acl['OwnerId'],
+            "Id": network_acl_id,
+            "Arn": f"arn:aws:ec2:{region}:{current_aws_account_id}:network-acl/{network_acl_id}",
+            "IsDefault": network_acl["IsDefault"],
+            "VpcId": network_acl["VpcId"],
+            "OwnerId": network_acl["OwnerId"],
         }
-        if network_acl.get('Associations') and network_acl['Associations']:
+        if network_acl.get("Associations") and network_acl["Associations"]:
             # Include subnet associations in the data object if they exist
-            for association in network_acl['Associations']:
-                base_network_acl['NetworkAclAssociationId'] = association['NetworkAclAssociationId']
-                base_network_acl['SubnetId'] = association['SubnetId']
+            for association in network_acl["Associations"]:
+                base_network_acl["NetworkAclAssociationId"] = association[
+                    "NetworkAclAssociationId"
+                ]
+                base_network_acl["SubnetId"] = association["SubnetId"]
                 network_acls.append(base_network_acl)
         else:
             # Otherwise if there's no associations then don't include that in the data object
@@ -66,21 +77,21 @@ def transform_network_acl_data(
 
         if network_acl.get("Entries"):
             for rule in network_acl["Entries"]:
-                direction = 'egress' if rule['Egress'] else 'inbound'
+                direction = "egress" if rule["Egress"] else "inbound"
                 transformed_rule = {
-                    'Id': f"{network_acl['NetworkAclId']}/{direction}/{rule['RuleNumber']}",
-                    'CidrBlock': rule.get('CidrBlock'),
-                    'Ipv6CidrBlock': rule.get('Ipv6CidrBlock'),
-                    'Egress': rule['Egress'],
-                    'Protocol': rule['Protocol'],
-                    'RuleAction': rule['RuleAction'],
-                    'RuleNumber': rule['RuleNumber'],
+                    "Id": f"{network_acl['NetworkAclId']}/{direction}/{rule['RuleNumber']}",
+                    "CidrBlock": rule.get("CidrBlock"),
+                    "Ipv6CidrBlock": rule.get("Ipv6CidrBlock"),
+                    "Egress": rule["Egress"],
+                    "Protocol": rule["Protocol"],
+                    "RuleAction": rule["RuleAction"],
+                    "RuleNumber": rule["RuleNumber"],
                     # Add pointer back to the nacl to create an edge
-                    'NetworkAclId': network_acl_id,
-                    'FromPort': rule.get('PortRange', {}).get('FromPort'),
-                    'ToPort': rule.get('PortRange', {}).get('ToPort'),
+                    "NetworkAclId": network_acl_id,
+                    "FromPort": rule.get("PortRange", {}).get("FromPort"),
+                    "ToPort": rule.get("PortRange", {}).get("ToPort"),
                 }
-                if transformed_rule['Egress']:
+                if transformed_rule["Egress"]:
                     outbound_rules.append(transformed_rule)
                 else:
                     inbound_rules.append(transformed_rule)
@@ -93,11 +104,11 @@ def transform_network_acl_data(
 
 @timeit
 def load_all_nacl_data(
-        neo4j_session: neo4j.Session,
-        ec2_acl_objects: Ec2AclObjects,
-        region: str,
-        aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    ec2_acl_objects: Ec2AclObjects,
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     load_network_acls(
         neo4j_session,
@@ -124,11 +135,11 @@ def load_all_nacl_data(
 
 @timeit
 def load_network_acls(
-        neo4j_session: neo4j.Session,
-        data: list[dict[str, Any]],
-        region: str,
-        aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     logger.info(f"Loading {len(data)} network acls in {region}.")
     load(
@@ -143,11 +154,11 @@ def load_network_acls(
 
 @timeit
 def load_network_acl_inbound_rules(
-        neo4j_session: neo4j.Session,
-        data: list[dict[str, Any]],
-        region: str,
-        aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     logger.info(f"Loading {len(data)} network acl inbound rules in {region}.")
     load(
@@ -162,11 +173,11 @@ def load_network_acl_inbound_rules(
 
 @timeit
 def load_network_acl_egress_rules(
-        neo4j_session: neo4j.Session,
-        data: list[dict[str, Any]],
-        region: str,
-        aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    region: str,
+    aws_account_id: str,
+    update_tag: int,
 ) -> None:
     logger.info(f"Loading {len(data)} network acl egress rules in {region}.")
     load(
@@ -180,23 +191,36 @@ def load_network_acl_egress_rules(
 
 
 @timeit
-def cleanup_network_acls(neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]) -> None:
-    GraphJob.from_node_schema(EC2NetworkAclSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(EC2NetworkAclInboundRuleSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(EC2NetworkAclEgressRuleSchema(), common_job_parameters).run(neo4j_session)
+def cleanup_network_acls(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    GraphJob.from_node_schema(EC2NetworkAclSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    GraphJob.from_node_schema(
+        EC2NetworkAclInboundRuleSchema(),
+        common_job_parameters,
+    ).run(neo4j_session)
+    GraphJob.from_node_schema(
+        EC2NetworkAclEgressRuleSchema(),
+        common_job_parameters,
+    ).run(neo4j_session)
 
 
 @timeit
 def sync_network_acls(
-        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 EC2 network ACLs for region '{region}' in account '{current_aws_account_id}'.")
+        logger.info(
+            f"Syncing EC2 network ACLs for region '{region}' in account '{current_aws_account_id}'.",
+        )
         data = get_network_acl_data(boto3_session, region)
         ec2_acl_data = transform_network_acl_data(data, region, current_aws_account_id)
         load_all_nacl_data(