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

cartography/intel/aws/ec2/auto_scaling_groups.py

@@ -1,243 +1,308 @@
 import logging
-from typing import Dict
-from typing import List
+from collections import namedtuple
+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 AutoScalingGroupSchema
+from cartography.models.aws.ec2.auto_scaling_groups import (
+    EC2InstanceAutoScalingGroupSchema,
+)
+from cartography.models.aws.ec2.auto_scaling_groups import (
+    EC2SubnetAutoScalingGroupSchema,
+)
+from cartography.models.aws.ec2.launch_configurations import LaunchConfigurationSchema
 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__)
 
+AsgData = namedtuple(
+    "AsgData",
+    [
+        "group_list",
+        "instance_list",
+        "subnet_list",
+    ],
+)
+
 
 @timeit
 @aws_handle_regions
-def get_ec2_auto_scaling_groups(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('autoscaling', region_name=region, config=get_botocore_config())
-    paginator = client.get_paginator('describe_auto_scaling_groups')
-    asgs: List[Dict] = []
+def get_ec2_auto_scaling_groups(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> list[dict]:
+    client = boto3_session.client(
+        "autoscaling",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    paginator = client.get_paginator("describe_auto_scaling_groups")
+    asgs: list[dict] = []
     for page in paginator.paginate():
-        asgs.extend(page['AutoScalingGroups'])
+        asgs.extend(page["AutoScalingGroups"])
     return asgs
 
 
 @timeit
 @aws_handle_regions
-def get_launch_configurations(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('autoscaling', region_name=region, config=get_botocore_config())
-    paginator = client.get_paginator('describe_launch_configurations')
-    lcs: List[Dict] = []
+def get_launch_configurations(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> list[dict]:
+    client = boto3_session.client(
+        "autoscaling",
+        region_name=region,
+        config=get_botocore_config(),
+    )
+    paginator = client.get_paginator("describe_launch_configurations")
+    lcs: list[dict] = []
     for page in paginator.paginate():
-        lcs.extend(page['LaunchConfigurations'])
+        lcs.extend(page["LaunchConfigurations"])
     return lcs
 
 
+def transform_launch_configurations(
+    configurations: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    transformed_configurations = []
+    for config in configurations:
+        transformed_configurations.append(
+            {
+                "AssociatePublicIpAddress": config.get("AssociatePublicIpAddress"),
+                "LaunchConfigurationARN": config.get("LaunchConfigurationARN"),
+                "LaunchConfigurationName": config.get("LaunchConfigurationName"),
+                "CreatedTime": config.get("CreatedTime"),
+                "ImageId": config.get("ImageId"),
+                "KeyName": config.get("KeyName"),
+                "SecurityGroups": config.get("SecurityGroups"),
+                "InstanceType": config.get("InstanceType"),
+                "KernelId": config.get("KernelId"),
+                "RamdiskId": config.get("RamdiskId"),
+                "InstanceMonitoring": config.get("InstanceMonitoring", {}).get(
+                    "Enabled",
+                ),
+                "SpotPrice": config.get("SpotPrice"),
+                "IamInstanceProfile": config.get("IamInstanceProfile"),
+                "EbsOptimized": config.get("EbsOptimized"),
+                "PlacementTenancy": config.get("PlacementTenancy"),
+            },
+        )
+    return transformed_configurations
+
+
+def transform_auto_scaling_groups(groups: list[dict[str, Any]]) -> AsgData:
+    transformed_groups = []
+    related_vpcs = []
+    related_instances = []
+    for group in groups:
+        transformed_groups.append(
+            {
+                "AutoScalingGroupARN": group["AutoScalingGroupARN"],
+                "CapacityRebalance": group.get("CapacityRebalance"),
+                "CreatedTime": str(group.get("CreatedTime")),
+                "DefaultCooldown": group.get("DefaultCooldown"),
+                "DesiredCapacity": group.get("DesiredCapacity"),
+                "HealthCheckGracePeriod": group.get("HealthCheckGracePeriod"),
+                "HealthCheckType": group.get("HealthCheckType"),
+                "LaunchConfigurationName": group.get("LaunchConfigurationName"),
+                "LaunchTemplateName": group.get("LaunchTemplate", {}).get(
+                    "LaunchTemplateName",
+                ),
+                "LaunchTemplateId": group.get("LaunchTemplate", {}).get(
+                    "LaunchTemplateId",
+                ),
+                "LaunchTemplateVersion": group.get("LaunchTemplate", {}).get("Version"),
+                "MaxInstanceLifetime": group.get("MaxInstanceLifetime"),
+                "MaxSize": group.get("MaxSize"),
+                "MinSize": group.get("MinSize"),
+                "AutoScalingGroupName": group.get("AutoScalingGroupName"),
+                "NewInstancesProtectedFromScaleIn": group.get(
+                    "NewInstancesProtectedFromScaleIn",
+                ),
+                "Status": group.get("Status"),
+            },
+        )
+
+        if group.get("VPCZoneIdentifier", None):
+            vpclist = group["VPCZoneIdentifier"]
+            subnet_ids = vpclist.split(",") if "," in vpclist else [vpclist]
+            subnets = []
+            for subnet_id in subnet_ids:
+                subnets.append(
+                    {
+                        "VPCZoneIdentifier": subnet_id,
+                        "AutoScalingGroupARN": group["AutoScalingGroupARN"],
+                    },
+                )
+            related_vpcs.extend(subnets)
+
+        for instance_data in group.get("Instances", []):
+            related_instances.append(
+                {
+                    "InstanceId": instance_data["InstanceId"],
+                    "AutoScalingGroupARN": group["AutoScalingGroupARN"],
+                },
+            )
+
+    return AsgData(
+        group_list=transformed_groups,
+        instance_list=related_instances,
+        subnet_list=related_vpcs,
+    )
+
+
 @timeit
 def load_launch_configurations(
-        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_lc = """
-    UNWIND $launch_configurations as lc
-        MERGE (config:LaunchConfiguration{id: lc.LaunchConfigurationARN})
-        ON CREATE SET config.firstseen = timestamp(), config.name = lc.LaunchConfigurationName,
-        config.arn = lc.LaunchConfigurationARN,
-        config.created_time = lc.CreatedTime
-        SET config.lastupdated = $update_tag, config.image_id = lc.ImageId,
-        config.key_name = lc.KeyName,
-        config.security_groups = lc.SecurityGroups,
-        config.instance_type = lc.InstanceType,
-        config.kernel_id = lc.KernelId,
-        config.ramdisk_id = lc.RamdiskId,
-        config.instance_monitoring_enabled = lc.InstanceMonitoring.Enabled,
-        config.spot_price = lc.SpotPrice,
-        config.iam_instance_profile = lc.IamInstanceProfile,
-        config.ebs_optimized = lc.EbsOptimized,
-        config.associate_public_ip_address = lc.AssociatePublicIpAddress,
-        config.placement_tenancy = lc.PlacementTenancy,
-        config.region=$Region
-        WITH config
-        MATCH (aa:AWSAccount{id: $AWS_ACCOUNT_ID})
-        MERGE (aa)-[r:RESOURCE]->(config)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-    """
-    for lc in data:
-        lc['CreatedTime'] = str(int(lc['CreatedTime'].timestamp()))
-
-    neo4j_session.run(
-        ingest_lc,
-        launch_configurations=data,
-        AWS_ACCOUNT_ID=current_aws_account_id,
+    load(
+        neo4j_session,
+        LaunchConfigurationSchema(),
+        data,
         Region=region,
-        update_tag=update_tag,
+        AWS_ID=current_aws_account_id,
+        lastupdated=update_tag,
     )
 
 
-@timeit
-def load_ec2_auto_scaling_groups(
-        neo4j_session: neo4j.Session, data: List[Dict], region: str, current_aws_account_id: str, update_tag: int,
+def load_groups(
+    neo4j_session: neo4j.Session,
+    data: list[dict],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
-    ingest_group = """
-    UNWIND $autoscaling_groups_list as ag
-        MERGE (group:AutoScalingGroup{arn: ag.AutoScalingGroupARN})
-        ON CREATE SET group.firstseen = timestamp(),
-        group.createdtime = ag.CreatedTime
-        SET group.launchconfigurationname = ag.LaunchConfigurationName,
-        group.launchtemplatename = ag.LaunchTemplate.LaunchTemplateName,
-        group.launchtemplateid = ag.LaunchTemplate.LaunchTemplateId,
-        group.launchtemplateversion = ag.LaunchTemplate.Version,
-        group.maxsize = ag.MaxSize, group.minsize = ag.MinSize, group.defaultcooldown = ag.DefaultCooldown,
-        group.desiredcapacity = ag.DesiredCapacity, group.healthchecktype = ag.HealthCheckType,
-        group.healthcheckgraceperiod = ag.HealthCheckGracePeriod, group.status = ag.Status,
-        group.newinstancesprotectedfromscalein = ag.NewInstancesProtectedFromScaleIn,
-        group.maxinstancelifetime = ag.MaxInstanceLifetime, group.capacityrebalance = ag.CapacityRebalance,
-        group.name = ag.AutoScalingGroupName,
-        group.lastupdated = $update_tag,
-        group.region=$Region
-        WITH group
-        MATCH (aa:AWSAccount{id: $AWS_ACCOUNT_ID})
-        MERGE (aa)-[r:RESOURCE]->(group)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-    """
-
-    ingest_vpc = """
-    UNWIND $vpc_id_list as vpc_id
-        MERGE (subnet:EC2Subnet{subnetid: vpc_id})
-        ON CREATE SET subnet.firstseen = timestamp()
-        SET subnet.lastupdated = $update_tag
-        WITH subnet
-        MATCH (group:AutoScalingGroup{arn: $GROUPARN})
-        MERGE (subnet)<-[r:VPC_IDENTIFIER]-(group)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-    """
-
-    ingest_instance = """
-    UNWIND $instances_list as i
-        MERGE (instance:Instance:EC2Instance{id: i.InstanceId})
-        ON CREATE SET instance.firstseen = timestamp()
-        SET instance.lastupdated = $update_tag, instance.region=$Region
-        WITH instance
-        MATCH (group:AutoScalingGroup{arn: $GROUPARN})
-        MERGE (instance)-[r:MEMBER_AUTO_SCALE_GROUP]->(group)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-        WITH instance
-        MATCH (aa:AWSAccount{id: $AWS_ACCOUNT_ID})
-        MERGE (aa)-[r:RESOURCE]->(instance)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-    """
-
-    ingest_lts = """
-    UNWIND $autoscaling_groups_list as ag
-        MATCH (group:AutoScalingGroup{arn: ag.AutoScalingGroupARN})
-        MATCH (template:LaunchTemplate{id: ag.LaunchTemplate.LaunchTemplateId})
-        MERGE (group)-[r:HAS_LAUNCH_TEMPLATE]->(template)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-    """
-
-    ingest_lcs = """
-    UNWIND $autoscaling_groups_list as ag
-        MATCH (group:AutoScalingGroup{arn: ag.AutoScalingGroupARN})
-        MATCH (config:LaunchConfiguration{name: ag.LaunchConfigurationName})
-        MERGE (group)-[r:HAS_LAUNCH_CONFIG]->(config)
-        ON CREATE SET r.firstseen = timestamp()
-        SET r.lastupdated = $update_tag
-    """
-
-    launch_configs = []
-    launch_templates = []
-    for group in data:
-        if group.get('LaunchConfigurationName'):
-            launch_configs.append(group)
-        if group.get('LaunchTemplate'):
-            launch_templates.append(group)
-
-        group['CreatedTime'] = str(group['CreatedTime'])
-
-    neo4j_session.run(
-        ingest_group,
-        autoscaling_groups_list=data,
-        AWS_ACCOUNT_ID=current_aws_account_id,
-        Region=region,
-        update_tag=update_tag,
-    )
-    neo4j_session.run(
-        ingest_lcs,
-        autoscaling_groups_list=launch_configs,
-        AWS_ACCOUNT_ID=current_aws_account_id,
+    load(
+        neo4j_session,
+        AutoScalingGroupSchema(),
+        data,
         Region=region,
-        update_tag=update_tag,
+        AWS_ID=current_aws_account_id,
+        lastupdated=update_tag,
     )
-    neo4j_session.run(
-        ingest_lts,
-        autoscaling_groups_list=launch_templates,
-        AWS_ACCOUNT_ID=current_aws_account_id,
+
+
+def load_asg_subnets(
+    neo4j_session: neo4j.Session,
+    data: list[dict],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        EC2SubnetAutoScalingGroupSchema(),
+        data,
         Region=region,
-        update_tag=update_tag,
+        AWS_ID=current_aws_account_id,
+        lastupdated=update_tag,
     )
 
-    for group in data:
-        group_arn = group["AutoScalingGroupARN"]
-        if group.get('VPCZoneIdentifier'):
-            vpclist = group["VPCZoneIdentifier"]
-            if ',' in vpclist:
-                data = vpclist.split(',')
-            else:
-                data = vpclist
-            neo4j_session.run(
-                ingest_vpc,
-                vpc_id_list=data,
-                GROUPARN=group_arn,
-                update_tag=update_tag,
-            )
 
-        if group.get("Instances"):
-            data = group["Instances"]
-            neo4j_session.run(
-                ingest_instance,
-                instances_list=data,
-                GROUPARN=group_arn,
-                AWS_ACCOUNT_ID=current_aws_account_id,
-                Region=region,
-                update_tag=update_tag,
-            )
+def load_asg_instances(
+    neo4j_session: neo4j.Session,
+    data: list[dict],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        EC2InstanceAutoScalingGroupSchema(),
+        data,
+        Region=region,
+        AWS_ID=current_aws_account_id,
+        lastupdated=update_tag,
+    )
 
 
 @timeit
-def cleanup_ec2_auto_scaling_groups(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job(
-        'aws_ingest_ec2_auto_scaling_groups_cleanup.json',
+def load_auto_scaling_groups(
+    neo4j_session: neo4j.Session,
+    data: AsgData,
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
+) -> None:
+    load_groups(
+        neo4j_session,
+        data.group_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_asg_instances(
         neo4j_session,
-        common_job_parameters,
+        data.instance_list,
+        region,
+        current_aws_account_id,
+        update_tag,
+    )
+    load_asg_subnets(
+        neo4j_session,
+        data.subnet_list,
+        region,
+        current_aws_account_id,
+        update_tag,
     )
 
 
 @timeit
-def cleanup_ec2_launch_configurations(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
-    run_cleanup_job(
-        'aws_import_ec2_launch_configurations_cleanup.json',
+def cleanup(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    logger.debug("Running EC2 instance cleanup")
+    GraphJob.from_node_schema(AutoScalingGroupSchema(), common_job_parameters).run(
+        neo4j_session,
+    )
+    GraphJob.from_node_schema(LaunchConfigurationSchema(), common_job_parameters).run(
         neo4j_session,
-        common_job_parameters,
     )
 
 
 @timeit
 def sync_ec2_auto_scaling_groups(
-        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 auto scaling groups for region '%s' in account '%s'.", region, current_aws_account_id)
+        logger.debug(
+            "Syncing auto scaling groups for region '%s' in account '%s'.",
+            region,
+            current_aws_account_id,
+        )
         lc_data = get_launch_configurations(boto3_session, region)
-        load_launch_configurations(neo4j_session, lc_data, region, current_aws_account_id, update_tag)
-        data = get_ec2_auto_scaling_groups(boto3_session, region)
-        load_ec2_auto_scaling_groups(neo4j_session, data, region, current_aws_account_id, update_tag)
-    cleanup_ec2_auto_scaling_groups(neo4j_session, common_job_parameters)
-    cleanup_ec2_launch_configurations(neo4j_session, common_job_parameters)
+        asg_data = get_ec2_auto_scaling_groups(boto3_session, region)
+        lc_transformed = transform_launch_configurations(lc_data)
+        asg_transformed = transform_auto_scaling_groups(asg_data)
+        load_launch_configurations(
+            neo4j_session,
+            lc_transformed,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
+        load_auto_scaling_groups(
+            neo4j_session,
+            asg_transformed,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
+    cleanup(neo4j_session, common_job_parameters)

cartography/intel/aws/ec2/elastic_ip_addresses.py

@@ -6,20 +6,29 @@ import boto3
 import neo4j
 from botocore.exceptions import ClientError
 
-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_elastic_ip_addresses(boto3_session: boto3.session.Session, region: str) -> List[Dict]:
-    client = boto3_session.client('ec2', region_name=region, config=get_botocore_config())
+def get_elastic_ip_addresses(
+    boto3_session: boto3.session.Session,
+    region: str,
+) -> List[Dict]:
+    client = boto3_session.client(
+        "ec2",
+        region_name=region,
+        config=get_botocore_config(),
+    )
     try:
-        addresses = client.describe_addresses()['Addresses']
+        addresses = client.describe_addresses()["Addresses"]
     except ClientError as e:
         logger.warning(f"Failed retrieve address for region - {region}. Error - {e}")
         raise
@@ -28,8 +37,11 @@ def get_elastic_ip_addresses(boto3_session: boto3.session.Session, region: str)
 
 @timeit
 def load_elastic_ip_addresses(
-    neo4j_session: neo4j.Session, elastic_ip_addresses: List[Dict], region: str,
-    current_aws_account_id: str, update_tag: int,
+    neo4j_session: neo4j.Session,
+    elastic_ip_addresses: List[Dict],
+    region: str,
+    current_aws_account_id: str,
+    update_tag: int,
 ) -> None:
     """
     Creates (:ElasticIpAddress)
@@ -37,7 +49,9 @@ def load_elastic_ip_addresses(
     (:EC2Instance)-[:ELASTIC_IP_ADDRESS]->(:ElasticIpAddress),
     (:NetworkInterface)-[:ELASTIC_IP_ADDRESS]->(:ElasticIpAddress),
     """
-    logger.info(f"Loading {len(elastic_ip_addresses)} Elastic IP Addresses in {region}.")
+    logger.info(
+        f"Loading {len(elastic_ip_addresses)} Elastic IP Addresses in {region}.",
+    )
     ingest_addresses = """
     UNWIND $elastic_ip_addresses as eia
         MERGE (address: ElasticIPAddress{id: eia.PublicIp})
@@ -70,7 +84,8 @@ def load_elastic_ip_addresses(
         SET r.lastupdated = $update_tag
     """
 
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingest_addresses,
         elastic_ip_addresses=elastic_ip_addresses,
         Region=region,
@@ -80,9 +95,12 @@ def load_elastic_ip_addresses(
 
 
 @timeit
-def cleanup_elastic_ip_addresses(neo4j_session: neo4j.Session, common_job_parameters: Dict) -> None:
+def cleanup_elastic_ip_addresses(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: Dict,
+) -> None:
     run_cleanup_job(
-        'aws_import_elastic_ip_addresses_cleanup.json',
+        "aws_import_elastic_ip_addresses_cleanup.json",
         neo4j_session,
         common_job_parameters,
     )
@@ -90,11 +108,23 @@ def cleanup_elastic_ip_addresses(neo4j_session: neo4j.Session, common_job_parame
 
 @timeit
 def sync_elastic_ip_addresses(
-    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(f"Syncing Elastic IP Addresses for region {region} in account {current_aws_account_id}.")
+        logger.info(
+            f"Syncing Elastic IP Addresses for region {region} in account {current_aws_account_id}.",
+        )
         addresses = get_elastic_ip_addresses(boto3_session, region)
-        load_elastic_ip_addresses(neo4j_session, addresses, region, current_aws_account_id, update_tag)
+        load_elastic_ip_addresses(
+            neo4j_session,
+            addresses,
+            region,
+            current_aws_account_id,
+            update_tag,
+        )
     cleanup_elastic_ip_addresses(neo4j_session, common_job_parameters)