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

cartography/intel/aws/ec2/subnets.py

@@ -5,31 +5,41 @@ 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.auto_scaling_groups import EC2SubnetAutoScalingGroupSchema
+from cartography.models.aws.ec2.auto_scaling_groups import (
+    EC2SubnetAutoScalingGroupSchema,
+)
 from cartography.models.aws.ec2.subnet_instance import EC2SubnetInstanceSchema
 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_subnet_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_subnets')
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    paginator = client.get_paginator("describe_subnets")
     subnets: List[Dict] = []
     for page in paginator.paginate():
-        subnets.extend(page['Subnets'])
+        subnets.extend(page["Subnets"])
     return subnets
 
 
 @timeit
 def load_subnets(
-    neo4j_session: neo4j.Session, data: List[Dict], region: str, aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    region: str,
+    aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
 
@@ -63,33 +73,59 @@ def load_subnets(
     """
 
     neo4j_session.run(
-        ingest_subnets, subnets=data, aws_update_tag=aws_update_tag,
-        region=region, aws_account_id=aws_account_id,
+        ingest_subnets,
+        subnets=data,
+        aws_update_tag=aws_update_tag,
+        region=region,
+        aws_account_id=aws_account_id,
     )
     neo4j_session.run(
-        ingest_subnet_vpc_relations, subnets=data, aws_update_tag=aws_update_tag,
-        region=region, aws_account_id=aws_account_id,
+        ingest_subnet_vpc_relations,
+        subnets=data,
+        aws_update_tag=aws_update_tag,
+        region=region,
+        aws_account_id=aws_account_id,
     )
     neo4j_session.run(
-        ingest_subnet_aws_account_relations, subnets=data, aws_update_tag=aws_update_tag,
-        region=region, aws_account_id=aws_account_id,
+        ingest_subnet_aws_account_relations,
+        subnets=data,
+        aws_update_tag=aws_update_tag,
+        region=region,
+        aws_account_id=aws_account_id,
     )
 
 
 @timeit
 def cleanup_subnets(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('aws_ingest_subnets_cleanup.json', neo4j_session, common_job_parameters)
-    GraphJob.from_node_schema(EC2SubnetInstanceSchema(), common_job_parameters).run(neo4j_session)
-    GraphJob.from_node_schema(EC2SubnetAutoScalingGroupSchema(), common_job_parameters).run(neo4j_session)
+    run_cleanup_job(
+        "aws_ingest_subnets_cleanup.json",
+        neo4j_session,
+        common_job_parameters,
+    )
+    GraphJob.from_node_schema(EC2SubnetInstanceSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    GraphJob.from_node_schema(
+        EC2SubnetAutoScalingGroupSchema(),
+        common_job_parameters,
+    ).run(neo4j_session)
 
 
 @timeit
 def sync_subnets(
-        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 subnets for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 subnets for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_subnet_data(boto3_session, region)
         load_subnets(neo4j_session, data, region, current_aws_account_id, update_tag)
     cleanup_subnets(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/tgw.py

@@ -6,18 +6,26 @@ import boto3
 import botocore.exceptions
 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_transit_gateways(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
+def get_transit_gateways(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
     data: List[Dict] = []
     try:
         data = client.describe_transit_gateways()["TransitGateways"]
@@ -25,51 +33,68 @@ def get_transit_gateways(boto3_session: boto3.session.Session, region: str) -> L
         # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/error-handling.html#parsing-error-responses-and-catching-exceptions-from-aws-services
         logger.warning(
             "Could not retrieve Transit Gateways due to boto3 error %s: %s. Skipping.",
-            e.response['Error']['Code'],
-            e.response['Error']['Message'],
+            e.response["Error"]["Code"],
+            e.response["Error"]["Message"],
         )
     return data
 
 
 @timeit
 @aws_handle_regions
-def get_tgw_attachments(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
+def get_tgw_attachments(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
     tgw_attachments: List[Dict] = []
     try:
-        paginator = client.get_paginator('describe_transit_gateway_attachments')
+        paginator = client.get_paginator("describe_transit_gateway_attachments")
         for page in paginator.paginate():
-            tgw_attachments.extend(page['TransitGatewayAttachments'])
+            tgw_attachments.extend(page["TransitGatewayAttachments"])
     except botocore.exceptions.ClientError as e:
         logger.warning(
             "Could not retrieve Transit Gateway Attachments due to boto3 error %s: %s. Skipping.",
-            e.response['Error']['Code'],
-            e.response['Error']['Message'],
+            e.response["Error"]["Code"],
+            e.response["Error"]["Message"],
         )
     return tgw_attachments
 
 
 @timeit
 @aws_handle_regions
-def get_tgw_vpc_attachments(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
+def get_tgw_vpc_attachments(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
     tgw_vpc_attachments: List[Dict] = []
     try:
-        paginator = client.get_paginator('describe_transit_gateway_vpc_attachments')
+        paginator = client.get_paginator("describe_transit_gateway_vpc_attachments")
         for page in paginator.paginate():
-            tgw_vpc_attachments.extend(page['TransitGatewayVpcAttachments'])
+            tgw_vpc_attachments.extend(page["TransitGatewayVpcAttachments"])
     except botocore.exceptions.ClientError as e:
         logger.warning(
             "Could not retrieve Transit Gateway VPC Attachments due to boto3 error %s: %s. Skipping.",
-            e.response['Error']['Code'],
-            e.response['Error']['Message'],
+            e.response["Error"]["Code"],
+            e.response["Error"]["Message"],
         )
     return tgw_vpc_attachments
 
 
 @timeit
 def load_transit_gateways(
-    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_transit_gateway = """
@@ -107,13 +132,21 @@ def load_transit_gateways(
             update_tag=update_tag,
         )
         _attach_shared_transit_gateway(
-            neo4j_session, tgw, region, current_aws_account_id, update_tag,
+            neo4j_session,
+            tgw,
+            region,
+            current_aws_account_id,
+            update_tag,
         )
 
 
 @timeit
 def _attach_shared_transit_gateway(
-    neo4j_session: neo4j.Session, tgw: Dict, region: str, current_aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    tgw: Dict,
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     attach_tgw = """
     MERGE (tgw:AWSTransitGateway {id: $ARN})
@@ -139,7 +172,10 @@ def _attach_shared_transit_gateway(
 
 @timeit
 def load_tgw_attachments(
-    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_transit_gateway = """
@@ -178,16 +214,25 @@ def load_tgw_attachments(
             update_tag=update_tag,
         )
 
-        if tgwa.get("VpcId"):  # only attach if the TGW attachment is a VPC TGW attachment
+        if tgwa.get(
+            "VpcId",
+        ):  # only attach if the TGW attachment is a VPC TGW attachment
             _attach_tgw_vpc_attachment_to_vpc_subnets(
-                neo4j_session, tgwa, region, current_aws_account_id, update_tag,
+                neo4j_session,
+                tgwa,
+                region,
+                current_aws_account_id,
+                update_tag,
             )
 
 
 @timeit
 def _attach_tgw_vpc_attachment_to_vpc_subnets(
-    neo4j_session: neo4j.Session, tgw_vpc_attachment: Dict, region: str,
-    current_aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    tgw_vpc_attachment: Dict,
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     """
     Attach a VPC Transit Gateway Attachment to the VPC and and subnets
@@ -233,27 +278,48 @@ def _attach_tgw_vpc_attachment_to_vpc_subnets(
 
 
 @timeit
-def cleanup_transit_gateways(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('aws_import_tgw_cleanup.json', neo4j_session, common_job_parameters)
+def cleanup_transit_gateways(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
+    run_cleanup_job("aws_import_tgw_cleanup.json", neo4j_session, common_job_parameters)
 
 
 def sync_transit_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 AWS Transit Gateways for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing AWS Transit Gateways for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         tgws = get_transit_gateways(boto3_session, region)
-        load_transit_gateways(neo4j_session, tgws, region, current_aws_account_id, update_tag)
+        load_transit_gateways(
+            neo4j_session,
+            tgws,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
 
         logger.debug(
             "Syncing AWS Transit Gateway Attachments for region '%s' in account '%s'.",
-            region, current_aws_account_id,
+            region,
+            current_aws_account_id,
         )
         tgw_attachments = get_tgw_attachments(boto3_session, region)
         tgw_vpc_attachments = get_tgw_vpc_attachments(boto3_session, region)
         load_tgw_attachments(
-            neo4j_session, tgw_attachments + tgw_vpc_attachments,
-            region, current_aws_account_id, update_tag,
+            neo4j_session,
+            tgw_attachments + tgw_vpc_attachments,
+            region,
+            current_aws_account_id,
+            update_tag,
         )
     cleanup_transit_gateways(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/util.py

@@ -6,6 +6,6 @@ def get_botocore_config() -> botocore.config.Config:
     return botocore.config.Config(
         read_timeout=360,
         retries={
-            'max_attempts': 10,
+            "max_attempts": 10,
         },
     )

cartography/intel/aws/ec2/volumes.py

@@ -18,38 +18,51 @@ logger = logging.getLogger(__name__)
 
 @timeit
 @aws_handle_regions
-def get_volumes(boto3_session: boto3.session.Session, region: str) -> List[Dict[str, Any]]:
-    client = boto3_session.client('ec2', region_name=region)
-    paginator = client.get_paginator('describe_volumes')
+def get_volumes(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict[str, Any]]:
+    client = boto3_session.client("ec2", region_name=region)
+    paginator = client.get_paginator("describe_volumes")
     volumes: List[Dict] = []
     for page in paginator.paginate():
-        volumes.extend(page['Volumes'])
+        volumes.extend(page["Volumes"])
     return volumes
 
 
-def transform_volumes(volumes: List[Dict[str, Any]], region: str, current_aws_account_id: str) -> List[Dict[str, Any]]:
+def transform_volumes(
+    volumes: List[Dict[str, Any]],
+    region: str,
+    current_aws_account_id: str,
+) -> List[Dict[str, Any]]:
     result = []
     for volume in volumes:
-        attachments = volume.get('Attachments', [])
-        active_attachments = [a for a in attachments if a['State'] == 'attached']
-
-        volume_id = volume['VolumeId']
-        raw_vol = ({
-            'Arn': build_arn('ec2', current_aws_account_id, 'volume', volume_id, region),
-            'AvailabilityZone': volume.get('AvailabilityZone'),
-            'CreateTime': volume.get('CreateTime'),
-            'Encrypted': volume.get('Encrypted'),
-            'Size': volume.get('Size'),
-            'State': volume.get('State'),
-            'OutpostArn': volume.get('OutpostArn'),
-            'SnapshotId': volume.get('SnapshotId'),
-            'Iops': volume.get('Iops'),
-            'FastRestored': volume.get('FastRestored'),
-            'MultiAttachEnabled': volume.get('MultiAttachEnabled'),
-            'VolumeType': volume.get('VolumeType'),
-            'VolumeId': volume_id,
-            'KmsKeyId': volume.get('KmsKeyId'),
-        })
+        attachments = volume.get("Attachments", [])
+        active_attachments = [a for a in attachments if a["State"] == "attached"]
+
+        volume_id = volume["VolumeId"]
+        raw_vol = {
+            "Arn": build_arn(
+                "ec2",
+                current_aws_account_id,
+                "volume",
+                volume_id,
+                region,
+            ),
+            "AvailabilityZone": volume.get("AvailabilityZone"),
+            "CreateTime": volume.get("CreateTime"),
+            "Encrypted": volume.get("Encrypted"),
+            "Size": volume.get("Size"),
+            "State": volume.get("State"),
+            "OutpostArn": volume.get("OutpostArn"),
+            "SnapshotId": volume.get("SnapshotId"),
+            "Iops": volume.get("Iops"),
+            "FastRestored": volume.get("FastRestored"),
+            "MultiAttachEnabled": volume.get("MultiAttachEnabled"),
+            "VolumeType": volume.get("VolumeType"),
+            "VolumeId": volume_id,
+            "KmsKeyId": volume.get("KmsKeyId"),
+        }
 
         if not active_attachments:
             result.append(raw_vol)
@@ -57,7 +70,7 @@ def transform_volumes(volumes: List[Dict[str, Any]], region: str, current_aws_ac
 
         for attachment in active_attachments:
             vol_with_attachment = raw_vol.copy()
-            vol_with_attachment['InstanceId'] = attachment['InstanceId']
+            vol_with_attachment["InstanceId"] = attachment["InstanceId"]
             result.append(vol_with_attachment)
 
     return result
@@ -65,11 +78,11 @@ def transform_volumes(volumes: List[Dict[str, Any]], region: str, current_aws_ac
 
 @timeit
 def load_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,
@@ -82,22 +95,37 @@ def load_volumes(
 
 
 @timeit
-def cleanup_volumes(neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]) -> None:
-    GraphJob.from_node_schema(EBSVolumeSchema(), common_job_parameters).run(neo4j_session)
+def cleanup_volumes(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict[str, Any],
+) -> None:
+    GraphJob.from_node_schema(EBSVolumeSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
 
 
 @timeit
 def sync_ebs_volumes(
-        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.debug("Syncing volumes for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.debug(
+            "Syncing volumes for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_volumes(boto3_session, region)
         transformed_data = transform_volumes(data, region, current_aws_account_id)
-        load_volumes(neo4j_session, transformed_data, region, current_aws_account_id, update_tag)
+        load_volumes(
+            neo4j_session,
+            transformed_data,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_volumes(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/vpc.py

@@ -6,23 +6,29 @@ 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_ec2_vpcs(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
-    return client.describe_vpcs()['Vpcs']
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    return client.describe_vpcs()["Vpcs"]
 
 
 def _get_cidr_association_statement(block_type: str) -> str:
-    INGEST_CIDR_TEMPLATE = Template("""
+    INGEST_CIDR_TEMPLATE = Template(
+        """
     MATCH (vpc:AWSVpc{id: $VpcId})
     WITH vpc
     UNWIND $CidrBlock as block_data
@@ -36,7 +42,8 @@ def _get_cidr_association_statement(block_type: str) -> str:
         WITH vpc, new_block
         MERGE (vpc)-[r:BLOCK_ASSOCIATION]->(new_block)
         ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag""")
+        SET r.lastupdated = $update_tag""",
+    )
 
     BLOCK_CIDR = "CidrBlock"
     STATE_NAME = "CidrBlockState"
@@ -53,12 +60,19 @@ def _get_cidr_association_statement(block_type: str) -> str:
     else:
         raise ValueError(f"Unsupported block type specified - {block_type}")
 
-    return INGEST_CIDR_TEMPLATE.safe_substitute(block_label=BLOCK_TYPE, block_cidr=BLOCK_CIDR, state_name=STATE_NAME)
+    return INGEST_CIDR_TEMPLATE.safe_substitute(
+        block_label=BLOCK_TYPE,
+        block_cidr=BLOCK_CIDR,
+        state_name=STATE_NAME,
+    )
 
 
 @timeit
 def load_cidr_association_set(
-    neo4j_session: neo4j.Session, vpc_id: str, vpc_data: Dict, block_type: str,
+    neo4j_session: neo4j.Session,
+    vpc_id: str,
+    vpc_data: Dict,
+    block_type: str,
     update_tag: int,
 ) -> None:
     ingest_statement = _get_cidr_association_statement(block_type)
@@ -78,7 +92,10 @@ def load_cidr_association_set(
 
 @timeit
 def load_ec2_vpcs(
-    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:
     # https://docs.aws.amazon.com/cli/latest/reference/ec2/describe-vpcs.html
@@ -161,16 +178,24 @@ def load_ec2_vpcs(
 
 @timeit
 def cleanup_ec2_vpcs(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job('aws_import_vpc_cleanup.json', neo4j_session, common_job_parameters)
+    run_cleanup_job("aws_import_vpc_cleanup.json", neo4j_session, common_job_parameters)
 
 
 @timeit
 def sync_vpc(
-    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 VPC for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.info(
+            "Syncing EC2 VPC for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         data = get_ec2_vpcs(boto3_session, region)
         load_ec2_vpcs(neo4j_session, data, region, current_aws_account_id, update_tag)
     cleanup_ec2_vpcs(neo4j_session, common_job_parameters)