cartography 0.93.0rc1__py3-none-any.whl → 0.123.0__py3-none-any.whl

cartography/intel/aws/ec2/subnets.py

@@ -1,93 +1,105 @@
 import logging
-from typing import Dict
-from typing import List
+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.auto_scaling_groups import (
+    EC2SubnetAutoScalingGroupSchema,
+)
 from cartography.models.aws.ec2.subnet_instance import EC2SubnetInstanceSchema
+from cartography.models.aws.ec2.subnets import EC2SubnetSchema
 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')
-    subnets: List[Dict] = []
+def get_subnet_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_subnets")
+    subnets: list[dict[str, Any]] = []
     for page in paginator.paginate():
-        subnets.extend(page['Subnets'])
+        subnets.extend(page["Subnets"])
     return subnets
 
 
+def transform_subnet_data(subnets: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Transform subnet data into a loadable format."""
+    transformed: list[dict[str, Any]] = []
+    for subnet in subnets:
+        transformed.append(subnet.copy())
+    return transformed
+
+
 @timeit
 def load_subnets(
-    neo4j_session: neo4j.Session, data: List[Dict], region: str, aws_account_id: str,
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    region: str,
+    aws_account_id: str,
     aws_update_tag: int,
 ) -> None:
-
-    ingest_subnets = """
-    UNWIND $subnets as subnet
-    MERGE (snet:EC2Subnet{subnetid: subnet.SubnetId})
-    ON CREATE SET snet.firstseen = timestamp()
-    SET snet.lastupdated = $aws_update_tag, snet.name = subnet.CidrBlock, snet.cidr_block = subnet.CidrBlock,
-    snet.available_ip_address_count = subnet.AvailableIpAddressCount, snet.default_for_az = subnet.DefaultForAz,
-    snet.map_customer_owned_ip_on_launch = subnet.MapCustomerOwnedIpOnLaunch,
-    snet.state = subnet.State, snet.assignipv6addressoncreation = subnet.AssignIpv6AddressOnCreation,
-    snet.map_public_ip_on_launch = subnet.MapPublicIpOnLaunch, snet.subnet_arn = subnet.SubnetArn,
-    snet.availability_zone = subnet.AvailabilityZone, snet.availability_zone_id = subnet.AvailabilityZoneId,
-    snet.subnetid = subnet.SubnetId
-    """
-
-    ingest_subnet_vpc_relations = """
-    UNWIND $subnets as subnet
-    MATCH (snet:EC2Subnet{subnetid: subnet.SubnetId}), (vpc:AWSVpc{id: subnet.VpcId})
-    MERGE (snet)-[r:MEMBER_OF_AWS_VPC]->(vpc)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $aws_update_tag
-    """
-
-    ingest_subnet_aws_account_relations = """
-    UNWIND $subnets as subnet
-    MATCH (snet:EC2Subnet{subnetid: subnet.SubnetId}), (aws:AWSAccount{id: $aws_account_id})
-    MERGE (aws)-[r:RESOURCE]->(snet)
-    ON CREATE SET r.firstseen = timestamp()
-    SET r.lastupdated = $aws_update_tag
-    """
-
-    neo4j_session.run(
-        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,
-    )
-    neo4j_session.run(
-        ingest_subnet_aws_account_relations, subnets=data, aws_update_tag=aws_update_tag,
-        region=region, aws_account_id=aws_account_id,
+    load(
+        neo4j_session,
+        EC2SubnetSchema(),
+        data,
+        Region=region,
+        AWS_ID=aws_account_id,
+        lastupdated=aws_update_tag,
     )
 
 
 @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)
+def cleanup_subnets(
+    neo4j_session: neo4j.Session, common_job_parameters: dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(EC2SubnetSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    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[str, Any],
 ) -> 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)
+        transformed = transform_subnet_data(data)
+        load_subnets(
+            neo4j_session,
+            transformed,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_subnets(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/tgw.py

@@ -6,18 +6,27 @@ import boto3
 import botocore.exceptions
 import neo4j
 
-from .util import get_botocore_config
+from cartography.client.core.tx import run_write_query
 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 +34,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 = """
@@ -95,7 +121,8 @@ def load_transit_gateways(
     for tgw in data:
         tgw_id = tgw["TransitGatewayId"]
 
-        neo4j_session.run(
+        run_write_query(
+            neo4j_session,
             ingest_transit_gateway,
             TgwId=tgw_id,
             ARN=tgw["TransitGatewayArn"],
@@ -107,13 +134,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})
@@ -128,7 +163,8 @@ def _attach_shared_transit_gateway(
     """
 
     if tgw["OwnerId"] != current_aws_account_id:
-        neo4j_session.run(
+        run_write_query(
+            neo4j_session,
             attach_tgw,
             ARN=tgw["TransitGatewayArn"],
             TransitGatewayId=tgw["TransitGatewayId"],
@@ -139,7 +175,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 = """
@@ -166,7 +205,8 @@ def load_tgw_attachments(
     for tgwa in data:
         tgwa_id = tgwa["TransitGatewayAttachmentId"]
 
-        neo4j_session.run(
+        run_write_query(
+            neo4j_session,
             ingest_transit_gateway,
             TgwAttachmentId=tgwa_id,
             TransitGatewayId=tgwa["TransitGatewayId"],
@@ -178,16 +218,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
@@ -216,7 +265,8 @@ def _attach_tgw_vpc_attachment_to_vpc_subnets(
     SET p.lastupdated = $update_tag
     """
 
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         attach_vpc_tgw_attachment_to_vpc,
         VpcId=tgw_vpc_attachment["VpcId"],
         TgwAttachmentId=tgw_vpc_attachment["TransitGatewayAttachmentId"],
@@ -224,7 +274,8 @@ def _attach_tgw_vpc_attachment_to_vpc_subnets(
     )
 
     for subnet_id in tgw_vpc_attachment["SubnetIds"]:
-        neo4j_session.run(
+        run_write_query(
+            neo4j_session,
             attach_vpc_tgw_attachment_to_subnet,
             SubnetId=subnet_id,
             TgwAttachmentId=tgw_vpc_attachment["TransitGatewayAttachmentId"],
@@ -233,27 +284,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.get("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)