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

cartography/intel/aws/ec2/reserved_instances.py

@@ -6,30 +6,43 @@ import boto3
 import neo4j
 from botocore.exceptions import ClientError
 
-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_reserved_instances(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
+def get_reserved_instances(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
     try:
-        reserved_instances = client.describe_reserved_instances()['ReservedInstances']
+        reserved_instances = client.describe_reserved_instances()["ReservedInstances"]
     except ClientError as e:
-        logger.warning(f"Failed retrieve reserved instances for region - {region}. Error - {e}")
+        logger.warning(
+            f"Failed retrieve reserved instances for region - {region}. Error - {e}",
+        )
         raise
     return reserved_instances
 
 
 @timeit
 def load_reserved_instances(
-        neo4j_session: neo4j.Session, data: List[Dict], region: str,
-        current_aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     ingest_reserved_instances = """
     UNWIND $reserved_instances_list as res
@@ -48,8 +61,8 @@ def load_reserved_instances(
     """
 
     for r_instance in data:
-        r_instance['Start'] = str(r_instance['Start'])
-        r_instance['End'] = str(r_instance['End'])
+        r_instance["Start"] = str(r_instance["Start"])
+        r_instance["End"] = str(r_instance["End"])
 
     neo4j_session.run(
         ingest_reserved_instances,
@@ -61,9 +74,12 @@ def load_reserved_instances(
 
 
 @timeit
-def cleanup_reserved_instances(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
+def cleanup_reserved_instances(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
     run_cleanup_job(
-        'aws_import_reserved_instances_cleanup.json',
+        "aws_import_reserved_instances_cleanup.json",
         neo4j_session,
         common_job_parameters,
     )
@@ -71,12 +87,25 @@ def cleanup_reserved_instances(neo4j_session: neo4j.Session, common_job_paramete
 
 @timeit
 def sync_ec2_reserved_instances(
-        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.debug("Syncing reserved instances for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.debug(
+            "Syncing reserved instances for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_reserved_instances(boto3_session, region)
-        load_reserved_instances(neo4j_session, data, region, current_aws_account_id, update_tag)
+        load_reserved_instances(
+            neo4j_session,
+            data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_reserved_instances(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/route_tables.py

@@ -7,7 +7,9 @@ import neo4j
 from cartography.client.core.tx import load
 from cartography.graph.job import GraphJob
 from cartography.intel.aws.ec2.util import get_botocore_config
-from cartography.models.aws.ec2.route_table_associations import RouteTableAssociationSchema
+from cartography.models.aws.ec2.route_table_associations import (
+    RouteTableAssociationSchema,
+)
 from cartography.models.aws.ec2.route_tables import RouteTableSchema
 from cartography.models.aws.ec2.routes import RouteSchema
 from cartography.util import aws_handle_regions
@@ -16,7 +18,9 @@ from cartography.util import timeit
 logger = logging.getLogger(__name__)
 
 
-def _get_route_id_and_target(route_table_id: str, route: dict[str, Any]) -> tuple[str, str | None]:
+def _get_route_id_and_target(
+    route_table_id: str, route: dict[str, Any]
+) -> tuple[str, str | None]:
     """
     Generate a unique identifier for an AWS EC2 route and return the target of the route
     regardless of its type.
@@ -29,18 +33,18 @@ def _get_route_id_and_target(route_table_id: str, route: dict[str, Any]) -> tupl
         A tuple containing the unique identifier for the route and the target of the route
     """
     route_target_keys = [
-        'DestinationCidrBlock',
-        'DestinationIpv6CidrBlock',
-        'GatewayId',
-        'InstanceId',
-        'NatGatewayId',
-        'TransitGatewayId',
-        'LocalGatewayId',
-        'CarrierGatewayId',
-        'NetworkInterfaceId',
-        'VpcPeeringConnectionId',
-        'EgressOnlyInternetGatewayId',
-        'CoreNetworkArn',
+        "DestinationCidrBlock",
+        "DestinationIpv6CidrBlock",
+        "GatewayId",
+        "InstanceId",
+        "NatGatewayId",
+        "TransitGatewayId",
+        "LocalGatewayId",
+        "CarrierGatewayId",
+        "NetworkInterfaceId",
+        "VpcPeeringConnectionId",
+        "EgressOnlyInternetGatewayId",
+        "CoreNetworkArn",
     ]
 
     # Start with the route table ID
@@ -63,23 +67,27 @@ def _get_route_id_and_target(route_table_id: str, route: dict[str, Any]) -> tupl
             "so that we can update the available keys.",
         )
 
-    return '|'.join(parts), target
+    return "|".join(parts), target
 
 
 @timeit
 @aws_handle_regions
-def get_route_tables(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_route_tables')
+def get_route_tables(
+    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_route_tables")
     route_tables: list[dict[str, Any]] = []
     for page in paginator.paginate():
-        route_tables.extend(page['RouteTables'])
+        route_tables.extend(page["RouteTables"])
     return route_tables
 
 
 def _transform_route_table_associations(
-        route_table_id: str,
-        associations: list[dict[str, Any]],
+    route_table_id: str,
+    associations: list[dict[str, Any]],
 ) -> tuple[list[dict[str, Any]], bool]:
     """
     Transform route table association data into a format suitable for cartography ingestion.
@@ -96,29 +104,33 @@ def _transform_route_table_associations(
     transformed = []
     is_main = False
     for association in associations:
-        if association.get('SubnetId'):
-            target = association['SubnetId']
-        elif association.get('GatewayId'):
-            target = association['GatewayId']
+        if association.get("SubnetId"):
+            target = association["SubnetId"]
+        elif association.get("GatewayId"):
+            target = association["GatewayId"]
         else:
             is_main = True
-            target = 'main'
+            target = "main"
 
         transformed_association = {
-            'id': association['RouteTableAssociationId'],
-            'route_table_id': route_table_id,
-            'subnet_id': association.get('SubnetId'),
-            'gateway_id': association.get('GatewayId'),
-            'main': association.get('Main', False),
-            'association_state': association.get('AssociationState', {}).get('State'),
-            'association_state_message': association.get('AssociationState', {}).get('Message'),
-            '_target': target,
+            "id": association["RouteTableAssociationId"],
+            "route_table_id": route_table_id,
+            "subnet_id": association.get("SubnetId"),
+            "gateway_id": association.get("GatewayId"),
+            "main": association.get("Main", False),
+            "association_state": association.get("AssociationState", {}).get("State"),
+            "association_state_message": association.get("AssociationState", {}).get(
+                "Message"
+            ),
+            "_target": target,
         }
         transformed.append(transformed_association)
     return transformed, is_main
 
 
-def _transform_route_table_routes(route_table_id: str, routes: list[dict[str, Any]]) -> list[dict[str, Any]]:
+def _transform_route_table_routes(
+    route_table_id: str, routes: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
     """
     Transform route table route data into a format suitable for cartography ingestion.
 
@@ -134,32 +146,32 @@ def _transform_route_table_routes(route_table_id: str, routes: list[dict[str, An
         route_id, target = _get_route_id_and_target(route_table_id, route)
 
         transformed_route = {
-            'id': route_id,
-            'route_table_id': route_table_id,
-            'destination_cidr_block': route.get('DestinationCidrBlock'),
-            'destination_ipv6_cidr_block': route.get('DestinationIpv6CidrBlock'),
-            'gateway_id': route.get('GatewayId'),
-            'instance_id': route.get('InstanceId'),
-            'instance_owner_id': route.get('InstanceOwnerId'),
-            'nat_gateway_id': route.get('NatGatewayId'),
-            'transit_gateway_id': route.get('TransitGatewayId'),
-            'local_gateway_id': route.get('LocalGatewayId'),
-            'carrier_gateway_id': route.get('CarrierGatewayId'),
-            'network_interface_id': route.get('NetworkInterfaceId'),
-            'vpc_peering_connection_id': route.get('VpcPeeringConnectionId'),
-            'state': route.get('State'),
-            'origin': route.get('Origin'),
-            'core_network_arn': route.get('CoreNetworkArn'),
-            'destination_prefix_list_id': route.get('DestinationPrefixListId'),
-            'egress_only_internet_gateway_id': route.get('EgressOnlyInternetGatewayId'),
-            '_target': target,
+            "id": route_id,
+            "route_table_id": route_table_id,
+            "destination_cidr_block": route.get("DestinationCidrBlock"),
+            "destination_ipv6_cidr_block": route.get("DestinationIpv6CidrBlock"),
+            "gateway_id": route.get("GatewayId"),
+            "instance_id": route.get("InstanceId"),
+            "instance_owner_id": route.get("InstanceOwnerId"),
+            "nat_gateway_id": route.get("NatGatewayId"),
+            "transit_gateway_id": route.get("TransitGatewayId"),
+            "local_gateway_id": route.get("LocalGatewayId"),
+            "carrier_gateway_id": route.get("CarrierGatewayId"),
+            "network_interface_id": route.get("NetworkInterfaceId"),
+            "vpc_peering_connection_id": route.get("VpcPeeringConnectionId"),
+            "state": route.get("State"),
+            "origin": route.get("Origin"),
+            "core_network_arn": route.get("CoreNetworkArn"),
+            "destination_prefix_list_id": route.get("DestinationPrefixListId"),
+            "egress_only_internet_gateway_id": route.get("EgressOnlyInternetGatewayId"),
+            "_target": target,
         }
         transformed.append(transformed_route)
     return transformed
 
 
 def transform_route_table_data(
-        route_tables: list[dict[str, Any]],
+    route_tables: list[dict[str, Any]],
 ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]:
     """
     Transform route table data into a format suitable for cartography ingestion.
@@ -175,31 +187,37 @@ def transform_route_table_data(
     route_data = []
 
     for rt in route_tables:
-        route_table_id = rt['RouteTableId']
+        route_table_id = rt["RouteTableId"]
 
         # Transform routes
         current_routes = []
-        if rt.get('Routes'):
-            current_routes = _transform_route_table_routes(route_table_id, rt['Routes'])
+        if rt.get("Routes"):
+            current_routes = _transform_route_table_routes(route_table_id, rt["Routes"])
             route_data.extend(current_routes)
 
         # If the rt has a association marked with main=True, then it is the main route table for the VPC.
         is_main = False
         # Transform associations
-        if rt.get('Associations'):
-            associations, is_main = _transform_route_table_associations(route_table_id, rt['Associations'])
+        if rt.get("Associations"):
+            associations, is_main = _transform_route_table_associations(
+                route_table_id, rt["Associations"]
+            )
             association_data.extend(associations)
 
         transformed_rt = {
-            'id': route_table_id,
-            'route_table_id': route_table_id,
-            'owner_id': rt.get('OwnerId'),
-            'vpc_id': rt.get('VpcId'),
-            'VpnGatewayIds': [vgw['GatewayId'] for vgw in rt.get('PropagatingVgws', [])],
-            'RouteTableAssociationIds': [assoc['RouteTableAssociationId'] for assoc in rt.get('Associations', [])],
-            'RouteIds': [route['id'] for route in current_routes],
-            'tags': rt.get('Tags', []),
-            'main': is_main,
+            "id": route_table_id,
+            "route_table_id": route_table_id,
+            "owner_id": rt.get("OwnerId"),
+            "vpc_id": rt.get("VpcId"),
+            "VpnGatewayIds": [
+                vgw["GatewayId"] for vgw in rt.get("PropagatingVgws", [])
+            ],
+            "RouteTableAssociationIds": [
+                assoc["RouteTableAssociationId"] for assoc in rt.get("Associations", [])
+            ],
+            "RouteIds": [route["id"] for route in current_routes],
+            "tags": rt.get("Tags", []),
+            "main": is_main,
         }
         transformed_tables.append(transformed_rt)
 
@@ -208,11 +226,11 @@ def transform_route_table_data(
 
 @timeit
 def load_route_tables(
-        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,
@@ -226,11 +244,11 @@ def load_route_tables(
 
 @timeit
 def load_route_table_associations(
-        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,
@@ -244,11 +262,11 @@ def load_route_table_associations(
 
 @timeit
 def load_routes(
-        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,
@@ -261,27 +279,49 @@ def load_routes(
 
 
 @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.debug("Running EC2 route tables cleanup")
-    GraphJob.from_node_schema(RouteTableSchema(), common_job_parameters).run(neo4j_session)
+    GraphJob.from_node_schema(RouteTableSchema(), common_job_parameters).run(
+        neo4j_session
+    )
     GraphJob.from_node_schema(RouteSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(RouteTableAssociationSchema(), common_job_parameters).run(neo4j_session)
+    GraphJob.from_node_schema(RouteTableAssociationSchema(), common_job_parameters).run(
+        neo4j_session
+    )
 
 
 @timeit
 def sync_route_tables(
-        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("Syncing EC2 route tables for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 route tables for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         route_tables = get_route_tables(boto3_session, region)
-        transformed_tables, association_data, route_data = transform_route_table_data(route_tables)
-        load_routes(neo4j_session, route_data, region, current_aws_account_id, update_tag)
-        load_route_table_associations(neo4j_session, association_data, region, current_aws_account_id, update_tag)
-        load_route_tables(neo4j_session, transformed_tables, region, current_aws_account_id, update_tag)
+        transformed_tables, association_data, route_data = transform_route_table_data(
+            route_tables
+        )
+        load_routes(
+            neo4j_session, route_data, region, current_aws_account_id, update_tag
+        )
+        load_route_table_associations(
+            neo4j_session, association_data, region, current_aws_account_id, update_tag
+        )
+        load_route_tables(
+            neo4j_session,
+            transformed_tables,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/security_groups.py

@@ -6,30 +6,46 @@ from typing import List
 import boto3
 import neo4j
 
-from .util import get_botocore_config
 from cartography.graph.job import GraphJob
-from cartography.models.aws.ec2.securitygroup_instance import EC2SecurityGroupInstanceSchema
+from cartography.models.aws.ec2.securitygroup_instance import (
+    EC2SecurityGroupInstanceSchema,
+)
 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_ec2_security_group_data(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    paginator = client.get_paginator('describe_security_groups')
+def get_ec2_security_group_data(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    paginator = client.get_paginator("describe_security_groups")
     security_groups: List[Dict] = []
     for page in paginator.paginate():
-        security_groups.extend(page['SecurityGroups'])
+        security_groups.extend(page["SecurityGroups"])
     return security_groups
 
 
 @timeit
-def load_ec2_security_group_rule(neo4j_session: neo4j.Session, group: Dict, rule_type: str, update_tag: int) -> None:
-    INGEST_RULE_TEMPLATE = Template("""
+def load_ec2_security_group_rule(
+    neo4j_session: neo4j.Session,
+    group: Dict,
+    rule_type: str,
+    update_tag: int,
+) -> None:
+    INGEST_RULE_TEMPLATE = Template(
+        """
     MERGE (rule:$rule_label{ruleid: $RuleId})
     ON CREATE SET rule :IpRule, rule.firstseen = timestamp(), rule.fromport = $FromPort, rule.toport = $ToPort,
     rule.protocol = $Protocol
@@ -39,7 +55,8 @@ def load_ec2_security_group_rule(neo4j_session: neo4j.Session, group: Dict, rule
     MERGE (group)<-[r:MEMBER_OF_EC2_SECURITY_GROUP]-(rule)
     ON CREATE SET r.firstseen = timestamp()
     SET r.lastupdated = $update_tag;
-    """)
+    """,
+    )
 
     ingest_rule_group_pair = """
     MERGE (group:EC2SecurityGroup{id: $GroupId})
@@ -64,7 +81,10 @@ def load_ec2_security_group_rule(neo4j_session: neo4j.Session, group: Dict, rule
     """
 
     group_id = group["GroupId"]
-    rule_type_map = {"IpPermissions": "IpPermissionInbound", "IpPermissionsEgress": "IpPermissionEgress"}
+    rule_type_map = {
+        "IpPermissions": "IpPermissionInbound",
+        "IpPermissionsEgress": "IpPermissionEgress",
+    }
 
     if group.get(rule_type):
         for rule in group[rule_type]:
@@ -76,7 +96,9 @@ def load_ec2_security_group_rule(neo4j_session: neo4j.Session, group: Dict, rule
             # NOTE Cypher query syntax is incompatible with Python string formatting, so we have to do this awkward
             # NOTE manual formatting instead.
             neo4j_session.run(
-                INGEST_RULE_TEMPLATE.safe_substitute(rule_label=rule_type_map[rule_type]),
+                INGEST_RULE_TEMPLATE.safe_substitute(
+                    rule_label=rule_type_map[rule_type],
+                ),
                 RuleId=ruleid,
                 FromPort=from_port,
                 ToPort=to_port,
@@ -104,8 +126,11 @@ def load_ec2_security_group_rule(neo4j_session: neo4j.Session, group: Dict, rule
 
 @timeit
 def load_ec2_security_groupinfo(
-    neo4j_session: neo4j.Session, data: List[Dict], region: str,
-    current_aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     ingest_security_group = """
     MERGE (group:EC2SecurityGroup{id: $GroupId})
@@ -138,26 +163,51 @@ def load_ec2_security_groupinfo(
         )
 
         load_ec2_security_group_rule(neo4j_session, group, "IpPermissions", update_tag)
-        load_ec2_security_group_rule(neo4j_session, group, "IpPermissionsEgress", update_tag)
+        load_ec2_security_group_rule(
+            neo4j_session,
+            group,
+            "IpPermissionsEgress",
+            update_tag,
+        )
 
 
 @timeit
-def cleanup_ec2_security_groupinfo(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
+def cleanup_ec2_security_groupinfo(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
     run_cleanup_job(
-        'aws_import_ec2_security_groupinfo_cleanup.json',
+        "aws_import_ec2_security_groupinfo_cleanup.json",
         neo4j_session,
         common_job_parameters,
     )
-    GraphJob.from_node_schema(EC2SecurityGroupInstanceSchema(), common_job_parameters).run(neo4j_session)
+    GraphJob.from_node_schema(
+        EC2SecurityGroupInstanceSchema(),
+        common_job_parameters,
+    ).run(neo4j_session)
 
 
 @timeit
 def sync_ec2_security_groupinfo(
-    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 security groups for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 security groups for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_ec2_security_group_data(boto3_session, region)
-        load_ec2_security_groupinfo(neo4j_session, data, region, current_aws_account_id, update_tag)
+        load_ec2_security_groupinfo(
+            neo4j_session,
+            data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_ec2_security_groupinfo(neo4j_session, common_job_parameters)