Fix subnet relationships.

cartography-cncf · Dec 4, 2024 · 9485df8 · 9485df8
1 parent f43be30
commit 9485df8
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 51 deletions.
diff --git 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
@@ -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
@@ -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,26 +191,13 @@ 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'
     properties: AutoScalingGroupNodeProperties = AutoScalingGroupNodeProperties()
     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
@@ -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"),
     }