Actually fix cartography-cncf#1210 using multiple classes

juju4 · Jul 14, 2023 · facb63b · facb63b
1 parent c0d9ac4
commit facb63b
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 3 deletions.
diff --git a/cartography/intel/aws/ec2/instances.py b/cartography/intel/aws/ec2/instances.py
@@ -17,7 +17,7 @@
 from cartography.models.aws.ec2.reservations import EC2ReservationSchema
 from cartography.models.aws.ec2.securitygroups import EC2SecurityGroupSchema
 from cartography.models.aws.ec2.subnets import EC2SubnetSchema
-from cartography.models.aws.ec2.volumes import EBSVolumeSchema
+from cartography.models.aws.ec2.volumes import EBSVolumeInstanceSchema
 from cartography.util import aws_handle_regions
 from cartography.util import timeit
 
@@ -273,7 +273,7 @@ def load_ec2_instance_ebs_volumes(
 ) -> None:
     load(
         neo4j_session,
-        EBSVolumeSchema(),
+        EBSVolumeInstanceSchema(),
         ebs_data,
         Region=region,
         AWS_ID=current_aws_account_id,

diff --git a/cartography/models/aws/ec2/volumes.py b/cartography/models/aws/ec2/volumes.py
@@ -18,7 +18,6 @@ class EBSVolumeNodeProperties(CartographyNodeProperties):
     volumeid: PropertyRef = PropertyRef('VolumeId', extra_index=True)
     region: PropertyRef = PropertyRef('Region', set_in_kwargs=True)
     lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True)
-    deleteontermination: PropertyRef = PropertyRef('DeleteOnTermination')
     availabilityzone: PropertyRef = PropertyRef('AvailabilityZone')
     createtime: PropertyRef = PropertyRef('CreateTime')
     encrypted: PropertyRef = PropertyRef('Encrypted')
@@ -67,6 +66,9 @@ class EBSVolumeToEC2Instance(CartographyRelSchema):
 
 @dataclass(frozen=True)
 class EBSVolumeSchema(CartographyNodeSchema):
+    """
+    EBS Volume properties as returned from the EBS Volume API response
+    """
     label: str = 'EBSVolume'
     properties: EBSVolumeNodeProperties = EBSVolumeNodeProperties()
     sub_resource_relationship: EBSVolumeToAWSAccount = EBSVolumeToAWSAccount()
@@ -75,3 +77,32 @@ class EBSVolumeSchema(CartographyNodeSchema):
             EBSVolumeToEC2Instance(),
         ],
     )
+
+
+@dataclass(frozen=True)
+class EBSVolumeInstanceProperties(CartographyNodeProperties):
+    """
+    EBS Volume properties as known by an EC2 instance.
+    The EC2 instance API response includes a `deleteontermination` field and a snapshot id.
+    """
+    arn: PropertyRef = PropertyRef('Arn', extra_index=True)
+    id: PropertyRef = PropertyRef('VolumeId')
+    volumeid: PropertyRef = PropertyRef('VolumeId', extra_index=True)
+    lastupdated: PropertyRef = PropertyRef('lastupdated', set_in_kwargs=True)
+    deleteontermination: PropertyRef = PropertyRef('DeleteOnTermination')
+    snapshotid: PropertyRef = PropertyRef('SnapshotId')
+
+
+@dataclass(frozen=True)
+class EBSVolumeInstanceSchema(CartographyNodeSchema):
+    """
+    EBS Volume from EC2 Instance API response. This is separate from `EBSVolumeSchema` to prevent issue #1210.
+    """
+    label: str = 'EBSVolume'
+    properties: EBSVolumeInstanceProperties = EBSVolumeInstanceProperties()
+    sub_resource_relationship: EBSVolumeToAWSAccount = EBSVolumeToAWSAccount()
+    other_relationships: OtherRelationships = OtherRelationships(
+        [
+            EBSVolumeToEC2Instance(),
+        ],
+    )