From 9485df88b31369a3d61dda4509e44f1a488daaf5 Mon Sep 17 00:00:00 2001 From: Sergio Franco Date: Wed, 4 Dec 2024 10:05:38 -0600 Subject: [PATCH] Fix subnet relationships. --- .../intel/aws/ec2/auto_scaling_groups.py | 9 +- cartography/intel/aws/ec2/subnets.py | 2 + .../models/aws/ec2/auto_scaling_groups.py | 82 +++++++++++-------- .../aws/ec2/test_ec2_auto_scaling_groups.py | 20 +---- 4 files changed, 62 insertions(+), 51 deletions(-) diff --git a/cartography/intel/aws/ec2/auto_scaling_groups.py b/cartography/intel/aws/ec2/auto_scaling_groups.py index 6bb5b67999..0cabbf6e1e 100644 --- a/cartography/intel/aws/ec2/auto_scaling_groups.py +++ b/cartography/intel/aws/ec2/auto_scaling_groups.py @@ -99,8 +99,13 @@ def transform_auto_scaling_groups(groups: list[dict[str, Any]]) -> AsgData: if group.get('VPCZoneIdentifier', None): vpclist = group['VPCZoneIdentifier'] subnet_ids = vpclist.split(',') if ',' in vpclist else [vpclist] - data = [{'VPCZoneIdentifier': subnet_id} for subnet_id in subnet_ids] - related_vpcs.extend(data) + 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({ diff --git a/cartography/intel/aws/ec2/subnets.py b/cartography/intel/aws/ec2/subnets.py index d306049835..a42aedbb4b 100644 --- a/cartography/intel/aws/ec2/subnets.py +++ b/cartography/intel/aws/ec2/subnets.py @@ -7,6 +7,7 @@ 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.subnet_instance import EC2SubnetInstanceSchema from cartography.util import aws_handle_regions from cartography.util import run_cleanup_job @@ -79,6 +80,7 @@ def load_subnets( 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) @timeit diff --git a/cartography/models/aws/ec2/auto_scaling_groups.py b/cartography/models/aws/ec2/auto_scaling_groups.py index 1cf2dc0334..c86fc6a707 100644 --- a/cartography/models/aws/ec2/auto_scaling_groups.py +++ b/cartography/models/aws/ec2/auto_scaling_groups.py @@ -1,6 +1,5 @@ from dataclasses import dataclass -from cartography.models.aws.ec2.subnet_instance import EC2SubnetToAWSAccount from cartography.models.core.common import PropertyRef from cartography.models.core.nodes import CartographyNodeProperties from cartography.models.core.nodes import CartographyNodeSchema @@ -36,13 +35,6 @@ class AutoScalingGroupNodeProperties(CartographyNodeProperties): lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) -@dataclass(frozen=True) -class EC2SubnetAutoScalingGroupNodeProperties(CartographyNodeProperties): - id: PropertyRef = PropertyRef('VPCZoneIdentifier') - subnetid: PropertyRef = PropertyRef('VPCZoneIdentifier', extra_index=True) - lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) - - # EC2 to AutoScalingGroup @dataclass(frozen=True) class EC2InstanceToAwsAccountRelProperties(CartographyRelProperties): @@ -96,35 +88,62 @@ class EC2InstanceAutoScalingGroupSchema(CartographyNodeSchema): ) +# EC2Subnet to AutoScalingGroup @dataclass(frozen=True) -class AutoScalingGroupToEC2SubnetRelProperties(CartographyRelProperties): +class EC2SubnetToAwsAccountRelProperties(CartographyRelProperties): lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) @dataclass(frozen=True) -class AutoScalingGroupToAwsAccountRelProperties(CartographyRelProperties): - lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) +class EC2SubnetToAWSAccount(CartographyRelSchema): + target_node_label: str = 'AWSAccount' + target_node_matcher: TargetNodeMatcher = make_target_node_matcher( + {'id': PropertyRef('AWS_ID', set_in_kwargs=True)}, + ) + direction: LinkDirection = LinkDirection.INWARD + rel_label: str = "RESOURCE" + properties: EC2SubnetToAwsAccountRelProperties = EC2SubnetToAwsAccountRelProperties() @dataclass(frozen=True) -class AutoScalingGroupToLaunchTemplateRelProperties(CartographyRelProperties): +class EC2SubnetToAutoScalingGroupRelProperties(CartographyRelProperties): lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) @dataclass(frozen=True) -class AutoScalingGroupToLaunchConfigurationRelProperties(CartographyRelProperties): +class EC2SubnetToAutoScalingGroup(CartographyRelSchema): + target_node_label: str = 'AutoScalingGroup' + target_node_matcher: TargetNodeMatcher = make_target_node_matcher( + {'id': PropertyRef('AutoScalingGroupARN')}, + ) + direction: LinkDirection = LinkDirection.INWARD + rel_label: str = "VPC_IDENTIFIER" + properties: EC2SubnetToAutoScalingGroupRelProperties = EC2SubnetToAutoScalingGroupRelProperties() + + +@dataclass(frozen=True) +class EC2SubnetAutoScalingGroupNodeProperties(CartographyNodeProperties): + id: PropertyRef = PropertyRef('VPCZoneIdentifier') + subnetid: PropertyRef = PropertyRef('VPCZoneIdentifier', extra_index=True) lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) @dataclass(frozen=True) -class AutoScalingGroupToEC2Subnet(CartographyRelSchema): - target_node_label: str = 'EC2Subnet' - target_node_matcher: TargetNodeMatcher = make_target_node_matcher( - {'id': PropertyRef('VPCZoneIdentifier')}, +class EC2SubnetAutoScalingGroupSchema(CartographyNodeSchema): + label: str = 'EC2Subnet' + properties: EC2SubnetAutoScalingGroupNodeProperties = EC2SubnetAutoScalingGroupNodeProperties() + sub_resource_relationship: EC2SubnetToAWSAccount = EC2SubnetToAWSAccount() + other_relationships: OtherRelationships = OtherRelationships( + [ + EC2SubnetToAutoScalingGroup(), + ], ) - direction: LinkDirection = LinkDirection.OUTWARD - rel_label: str = "VPC_IDENTIFIER" - properties: AutoScalingGroupToEC2SubnetRelProperties = AutoScalingGroupToEC2SubnetRelProperties() + + +# AutoScalingGroup +@dataclass(frozen=True) +class AutoScalingGroupToAwsAccountRelProperties(CartographyRelProperties): + lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) @dataclass(frozen=True) @@ -138,6 +157,11 @@ class AutoScalingGroupToAWSAccount(CartographyRelSchema): properties: AutoScalingGroupToAwsAccountRelProperties = AutoScalingGroupToAwsAccountRelProperties() +@dataclass(frozen=True) +class AutoScalingGroupToLaunchTemplateRelProperties(CartographyRelProperties): + lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) + + @dataclass(frozen=True) class AutoScalingGroupToLaunchTemplate(CartographyRelSchema): target_node_label: str = 'LaunchTemplate' @@ -149,6 +173,11 @@ class AutoScalingGroupToLaunchTemplate(CartographyRelSchema): properties: AutoScalingGroupToLaunchTemplateRelProperties = AutoScalingGroupToLaunchTemplateRelProperties() +@dataclass(frozen=True) +class AutoScalingGroupToLaunchConfigurationRelProperties(CartographyRelProperties): + lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True) + + @dataclass(frozen=True) class AutoScalingGroupToLaunchConfiguration(CartographyRelSchema): target_node_label: str = 'LaunchConfiguration' @@ -162,18 +191,6 @@ class AutoScalingGroupToLaunchConfiguration(CartographyRelSchema): ) -@dataclass(frozen=True) -class EC2SubnetAutoScalingGroupSchema(CartographyNodeSchema): - label: str = 'EC2Subnet' - properties: EC2SubnetAutoScalingGroupNodeProperties = EC2SubnetAutoScalingGroupNodeProperties() - sub_resource_relationship: EC2SubnetToAWSAccount = EC2SubnetToAWSAccount() - other_relationships: OtherRelationships = OtherRelationships( - [ - AutoScalingGroupToEC2Subnet(), - ], - ) - - @dataclass(frozen=True) class AutoScalingGroupSchema(CartographyNodeSchema): label: str = 'AutoScalingGroup' @@ -181,7 +198,6 @@ class AutoScalingGroupSchema(CartographyNodeSchema): sub_resource_relationship: AutoScalingGroupToAWSAccount = AutoScalingGroupToAWSAccount() other_relationships: OtherRelationships = OtherRelationships( [ - AutoScalingGroupToEC2Subnet(), AutoScalingGroupToLaunchTemplate(), AutoScalingGroupToLaunchConfiguration(), ], diff --git a/tests/integration/cartography/intel/aws/ec2/test_ec2_auto_scaling_groups.py b/tests/integration/cartography/intel/aws/ec2/test_ec2_auto_scaling_groups.py index c7eff6fc6c..fba2ec3ce0 100644 --- a/tests/integration/cartography/intel/aws/ec2/test_ec2_auto_scaling_groups.py +++ b/tests/integration/cartography/intel/aws/ec2/test_ec2_auto_scaling_groups.py @@ -90,20 +90,8 @@ def test_sync_ec2_auto_scaling_groups(mock_get_instances, mock_get_launch_config rel_label='MEMBER_AUTO_SCALE_GROUP', rel_direction_right=False, ) == { - ( - GET_AUTO_SCALING_GROUPS[1]['AutoScalingGroupARN'], - DESCRIBE_INSTANCES['Reservations'][0]['Instances'][0]['InstanceId'], - ), - ( - GET_AUTO_SCALING_GROUPS[1]['AutoScalingGroupARN'], - DESCRIBE_INSTANCES['Reservations'][1]['Instances'][0]['InstanceId'], - ), - ( - GET_AUTO_SCALING_GROUPS[1]['AutoScalingGroupARN'], - DESCRIBE_INSTANCES['Reservations'][2]['Instances'][0]['InstanceId'], - ), - ( - GET_AUTO_SCALING_GROUPS[1]['AutoScalingGroupARN'], - DESCRIBE_INSTANCES['Reservations'][2]['Instances'][1]['InstanceId'], - ), + (GET_AUTO_SCALING_GROUPS[1]['AutoScalingGroupARN'], "i-01"), + (GET_AUTO_SCALING_GROUPS[1]['AutoScalingGroupARN'], "i-02"), + (GET_AUTO_SCALING_GROUPS[1]['AutoScalingGroupARN'], "i-03"), + (GET_AUTO_SCALING_GROUPS[1]['AutoScalingGroupARN'], "i-04"), }