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

cartography/intel/aws/ec2/instances.py

@@ -11,12 +11,18 @@ 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.auto_scaling_groups import EC2InstanceAutoScalingGroupSchema
+from cartography.models.aws.ec2.auto_scaling_groups import (
+    EC2InstanceAutoScalingGroupSchema,
+)
 from cartography.models.aws.ec2.instances import EC2InstanceSchema
 from cartography.models.aws.ec2.keypair_instance import EC2KeyPairInstanceSchema
-from cartography.models.aws.ec2.networkinterface_instance import EC2NetworkInterfaceInstanceSchema
+from cartography.models.aws.ec2.networkinterface_instance import (
+    EC2NetworkInterfaceInstanceSchema,
+)
 from cartography.models.aws.ec2.reservations import EC2ReservationSchema
-from cartography.models.aws.ec2.securitygroup_instance import EC2SecurityGroupInstanceSchema
+from cartography.models.aws.ec2.securitygroup_instance import (
+    EC2SecurityGroupInstanceSchema,
+)
 from cartography.models.aws.ec2.subnet_instance import EC2SubnetInstanceSchema
 from cartography.models.aws.ec2.volumes import EBSVolumeInstanceSchema
 from cartography.util import aws_handle_regions
@@ -25,7 +31,8 @@ from cartography.util import timeit
 logger = logging.getLogger(__name__)
 
 Ec2Data = namedtuple(
-    'Ec2Data', [
+    "Ec2Data",
+    [
         "reservation_list",
         "instance_list",
         "subnet_list",
@@ -40,15 +47,23 @@ Ec2Data = namedtuple(
 @timeit
 @aws_handle_regions
 def get_ec2_instances(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_instances')
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    paginator = client.get_paginator("describe_instances")
     reservations: List[Dict[str, Any]] = []
     for page in paginator.paginate():
-        reservations.extend(page['Reservations'])
+        reservations.extend(page["Reservations"])
     return reservations
 
 
-def transform_ec2_instances(reservations: List[Dict[str, Any]], region: str, current_aws_account_id: str) -> Ec2Data:
+def transform_ec2_instances(
+    reservations: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+) -> Ec2Data:
     reservation_list = []
     instance_list = []
     subnet_list = []
@@ -58,92 +73,119 @@ def transform_ec2_instances(reservations: List[Dict[str, Any]], region: str, cur
     instance_ebs_volumes_list = []
 
     for reservation in reservations:
-        reservation_id = reservation['ReservationId']
-        reservation_list.append({
-            'RequesterId': reservation.get('RequesterId'),
-            'ReservationId': reservation['ReservationId'],
-            'OwnerId': reservation['OwnerId'],
-        })
-        for instance in reservation['Instances']:
-            instance_id = instance['InstanceId']
+        reservation_id = reservation["ReservationId"]
+        reservation_list.append(
+            {
+                "RequesterId": reservation.get("RequesterId"),
+                "ReservationId": reservation["ReservationId"],
+                "OwnerId": reservation["OwnerId"],
+            },
+        )
+        for instance in reservation["Instances"]:
+            instance_id = instance["InstanceId"]
             launch_time = instance.get("LaunchTime")
-            launch_time_unix = str(time.mktime(launch_time.timetuple())) if launch_time else None
+            launch_time_unix = (
+                str(time.mktime(launch_time.timetuple())) if launch_time else None
+            )
             instance_list.append(
                 {
-                    'InstanceId': instance_id,
-                    'ReservationId': reservation_id,
-                    'PublicDnsName': instance.get("PublicDnsName"),
-                    'PublicIpAddress': instance.get("PublicIpAddress"),
-                    'PrivateIpAddress': instance.get("PrivateIpAddress"),
-                    'ImageId': instance.get("ImageId"),
-                    'InstanceType': instance.get("InstanceType"),
-                    'IamInstanceProfile': instance.get("IamInstanceProfile", {}).get("Arn"),
-                    'MonitoringState': instance.get("Monitoring", {}).get("State"),
-                    'LaunchTime': instance.get("LaunchTime"),
-                    'LaunchTimeUnix': launch_time_unix,
-                    'State': instance.get("State", {}).get("Name"),
-                    'AvailabilityZone': instance.get("Placement", {}).get("AvailabilityZone"),
-                    'Tenancy': instance.get("Placement", {}).get("Tenancy"),
-                    'HostResourceGroupArn': instance.get("Placement", {}).get("HostResourceGroupArn"),
-                    'Platform': instance.get("Platform"),
-                    'Architecture': instance.get("Architecture"),
-                    'EbsOptimized': instance.get("EbsOptimized"),
-                    'BootMode': instance.get("BootMode"),
-                    'InstanceLifecycle': instance.get("InstanceLifecycle"),
-                    'HibernationOptions': instance.get("HibernationOptions", {}).get("Configured"),
+                    "InstanceId": instance_id,
+                    "ReservationId": reservation_id,
+                    "PublicDnsName": instance.get("PublicDnsName"),
+                    "PublicIpAddress": instance.get("PublicIpAddress"),
+                    "PrivateIpAddress": instance.get("PrivateIpAddress"),
+                    "ImageId": instance.get("ImageId"),
+                    "InstanceType": instance.get("InstanceType"),
+                    "IamInstanceProfile": instance.get("IamInstanceProfile", {}).get(
+                        "Arn",
+                    ),
+                    "MonitoringState": instance.get("Monitoring", {}).get("State"),
+                    "LaunchTime": instance.get("LaunchTime"),
+                    "LaunchTimeUnix": launch_time_unix,
+                    "State": instance.get("State", {}).get("Name"),
+                    "AvailabilityZone": instance.get("Placement", {}).get(
+                        "AvailabilityZone",
+                    ),
+                    "Tenancy": instance.get("Placement", {}).get("Tenancy"),
+                    "HostResourceGroupArn": instance.get("Placement", {}).get(
+                        "HostResourceGroupArn",
+                    ),
+                    "Platform": instance.get("Platform"),
+                    "Architecture": instance.get("Architecture"),
+                    "EbsOptimized": instance.get("EbsOptimized"),
+                    "BootMode": instance.get("BootMode"),
+                    "InstanceLifecycle": instance.get("InstanceLifecycle"),
+                    "HibernationOptions": instance.get("HibernationOptions", {}).get(
+                        "Configured",
+                    ),
                 },
             )
 
-            subnet_id = instance.get('SubnetId')
+            subnet_id = instance.get("SubnetId")
             if subnet_id:
                 subnet_list.append(
                     {
-                        'SubnetId': subnet_id,
-                        'InstanceId': instance_id,
+                        "SubnetId": subnet_id,
+                        "InstanceId": instance_id,
                     },
                 )
 
             if instance.get("KeyName"):
                 key_name = instance["KeyName"]
-                key_pair_arn = f'arn:aws:ec2:{region}:{current_aws_account_id}:key-pair/{key_name}'
-                keypair_list.append({
-                    'KeyPairArn': key_pair_arn,
-                    'KeyName': key_name,
-                    'InstanceId': instance_id,
-                })
+                key_pair_arn = (
+                    f"arn:aws:ec2:{region}:{current_aws_account_id}:key-pair/{key_name}"
+                )
+                keypair_list.append(
+                    {
+                        "KeyPairArn": key_pair_arn,
+                        "KeyName": key_name,
+                        "InstanceId": instance_id,
+                    },
+                )
 
             if instance.get("SecurityGroups"):
                 for group in instance["SecurityGroups"]:
                     sg_list.append(
                         {
-                            'GroupId': group['GroupId'],
-                            'InstanceId': instance_id,
+                            "GroupId": group["GroupId"],
+                            "InstanceId": instance_id,
                         },
                     )
 
-            for network_interface in instance['NetworkInterfaces']:
-                for security_group in network_interface.get('Groups', []):
-                    network_interface_list.append({
-                        'NetworkInterfaceId': network_interface['NetworkInterfaceId'],
-                        'Status': network_interface['Status'],
-                        'MacAddress': network_interface['MacAddress'],
-                        'Description': network_interface['Description'],
-                        'PrivateDnsName': network_interface.get('PrivateDnsName'),
-                        'PrivateIpAddress': network_interface['PrivateIpAddress'],
-                        'InstanceId': instance_id,
-                        'SubnetId': subnet_id,
-                        'GroupId': security_group['GroupId'],
-                    })
-
-            if 'BlockDeviceMappings' in instance and len(instance['BlockDeviceMappings']) > 0:
-                for mapping in instance['BlockDeviceMappings']:
-                    if 'VolumeId' in mapping['Ebs']:
-                        instance_ebs_volumes_list.append({
-                            'InstanceId': instance_id,
-                            'VolumeId': mapping['Ebs']['VolumeId'],
-                            'DeleteOnTermination': mapping['Ebs']['DeleteOnTermination'],
-                            # 'SnapshotId': mapping['Ebs']['SnapshotId'],  # TODO check on this
-                        })
+            for network_interface in instance["NetworkInterfaces"]:
+                for security_group in network_interface.get("Groups", []):
+                    network_interface_list.append(
+                        {
+                            "NetworkInterfaceId": network_interface[
+                                "NetworkInterfaceId"
+                            ],
+                            "Status": network_interface["Status"],
+                            "MacAddress": network_interface["MacAddress"],
+                            "Description": network_interface["Description"],
+                            "PrivateDnsName": network_interface.get("PrivateDnsName"),
+                            "PrivateIpAddress": network_interface["PrivateIpAddress"],
+                            "InstanceId": instance_id,
+                            "SubnetId": subnet_id,
+                            "GroupId": security_group["GroupId"],
+                        },
+                    )
+
+            if (
+                "BlockDeviceMappings" in instance
+                and len(instance["BlockDeviceMappings"]) > 0
+            ):
+                for mapping in instance["BlockDeviceMappings"]:
+                    if "VolumeId" in mapping["Ebs"]:
+                        instance_ebs_volumes_list.append(
+                            {
+                                "InstanceId": instance_id,
+                                "VolumeId": mapping["Ebs"]["VolumeId"],
+                                "DeleteOnTermination": mapping["Ebs"][
+                                    "DeleteOnTermination"
+                                ],
+                                # 'SnapshotId': mapping['Ebs']['SnapshotId'],  # TODO check on this
+                            },
+                        )
 
     return Ec2Data(
         reservation_list=reservation_list,
@@ -158,11 +200,11 @@ def transform_ec2_instances(reservations: List[Dict[str, Any]], region: str, cur
 
 @timeit
 def load_ec2_reservations(
-        neo4j_session: neo4j.Session,
-        reservation_list: List[Dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    reservation_list: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     load(
         neo4j_session,
@@ -176,11 +218,11 @@ def load_ec2_reservations(
 
 @timeit
 def load_ec2_subnets(
-        neo4j_session: neo4j.Session,
-        subnet_list: List[Dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    subnet_list: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     load(
         neo4j_session,
@@ -194,11 +236,11 @@ def load_ec2_subnets(
 
 @timeit
 def load_ec2_keypair_instances(
-        neo4j_session: neo4j.Session,
-        key_pair_list: List[Dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    key_pair_list: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     # Load EC2 keypairs as known by describe-instances.
     load(
@@ -213,11 +255,11 @@ def load_ec2_keypair_instances(
 
 @timeit
 def load_ec2_security_groups(
-        neo4j_session: neo4j.Session,
-        sg_list: List[Dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    sg_list: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     load(
         neo4j_session,
@@ -231,11 +273,11 @@ def load_ec2_security_groups(
 
 @timeit
 def load_ec2_network_interfaces(
-        neo4j_session: neo4j.Session,
-        network_interface_list: List[Dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    network_interface_list: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     load(
         neo4j_session,
@@ -249,11 +291,11 @@ def load_ec2_network_interfaces(
 
 @timeit
 def load_ec2_instance_nodes(
-        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:
     load(
         neo4j_session,
@@ -267,11 +309,11 @@ def load_ec2_instance_nodes(
 
 @timeit
 def load_ec2_instance_ebs_volumes(
-        neo4j_session: neo4j.Session,
-        ebs_data: List[Dict[str, Any]],
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
+    neo4j_session: neo4j.Session,
+    ebs_data: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     load(
         neo4j_session,
@@ -284,46 +326,102 @@ def load_ec2_instance_ebs_volumes(
 
 
 def load_ec2_instance_data(
-        neo4j_session: neo4j.Session,
-        region: str,
-        current_aws_account_id: str,
-        update_tag: int,
-        reservation_list: List[Dict[str, Any]],
-        instance_list: List[Dict[str, Any]],
-        subnet_list: List[Dict[str, Any]],
-        sg_list: List[Dict[str, Any]],
-        key_pair_list: List[Dict[str, Any]],
-        nic_list: List[Dict[str, Any]],
-        ebs_volumes_list: List[Dict[str, Any]],
+    neo4j_session: neo4j.Session,
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
+    reservation_list: List[Dict[str, Any]],
+    instance_list: List[Dict[str, Any]],
+    subnet_list: List[Dict[str, Any]],
+    sg_list: List[Dict[str, Any]],
+    key_pair_list: List[Dict[str, Any]],
+    nic_list: List[Dict[str, Any]],
+    ebs_volumes_list: List[Dict[str, Any]],
 ) -> None:
-    load_ec2_reservations(neo4j_session, reservation_list, region, current_aws_account_id, update_tag)
-    load_ec2_instance_nodes(neo4j_session, instance_list, region, current_aws_account_id, update_tag)
-    load_ec2_subnets(neo4j_session, subnet_list, region, current_aws_account_id, update_tag)
-    load_ec2_security_groups(neo4j_session, sg_list, region, current_aws_account_id, update_tag)
-    load_ec2_keypair_instances(neo4j_session, key_pair_list, region, current_aws_account_id, update_tag)
-    load_ec2_network_interfaces(neo4j_session, nic_list, region, current_aws_account_id, update_tag)
-    load_ec2_instance_ebs_volumes(neo4j_session, ebs_volumes_list, region, current_aws_account_id, update_tag)
+    load_ec2_reservations(
+        neo4j_session,
+        reservation_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_ec2_instance_nodes(
+        neo4j_session,
+        instance_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_ec2_subnets(
+        neo4j_session,
+        subnet_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_ec2_security_groups(
+        neo4j_session,
+        sg_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_ec2_keypair_instances(
+        neo4j_session,
+        key_pair_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_ec2_network_interfaces(
+        neo4j_session,
+        nic_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_ec2_instance_ebs_volumes(
+        neo4j_session,
+        ebs_volumes_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
 
 
 @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 instance cleanup")
-    GraphJob.from_node_schema(EC2ReservationSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(EC2InstanceSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(EC2InstanceAutoScalingGroupSchema(), common_job_parameters).run(neo4j_session)
+    GraphJob.from_node_schema(EC2ReservationSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    GraphJob.from_node_schema(EC2InstanceSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    GraphJob.from_node_schema(
+        EC2InstanceAutoScalingGroupSchema(),
+        common_job_parameters,
+    ).run(neo4j_session)
 
 
 @timeit
 def sync_ec2_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[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 instances for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 instances for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         reservations = get_ec2_instances(boto3_session, region)
         ec2_data = transform_ec2_instances(reservations, region, current_aws_account_id)
         load_ec2_instance_data(

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)