#996: AWS Fix KeyError exception when getting DNSName for load_balancer_id

cartography/intel/aws/ec2/load_balancer_v2s.py

@@ -77,14 +77,19 @@ def load_load_balancer_v2s(
     SET r.lastupdated = $update_tag
     """
     for lb in data:
-        load_balancer_id = lb["DNSName"]
+        # every load balancer has an arn that can be used as unique id instead of DNSName
+        # LoadBalancers V2 of type gateway do not contain a DNSName field
+        load_balancer_id = lb["LoadBalancerArn"]
+
+        # if a load balancer has dns name, it'll return the value else it won't set in Neo4j
+        dns_name = lb.get("DNSName", None)
 
         neo4j_session.run(
             ingest_load_balancer_v2,
             ID=load_balancer_id,
             CREATED_TIME=str(lb["CreatedTime"]),
             NAME=lb["LoadBalancerName"],
-            DNS_NAME=load_balancer_id,
+            DNS_NAME=dns_name,
             HOSTED_ZONE_NAME_ID=lb.get("CanonicalHostedZoneNameID"),
             ELBv2_TYPE=lb.get("Type"),
             SCHEME=lb.get("Scheme"),

cartography/intel/aws/ec2/network_interfaces.py

@@ -161,7 +161,7 @@ def load_network_interface_elbv2_relations(
     ingest_network_interface_elb2_relations = """
     UNWIND $elb_associations AS elb_association
         MATCH (netinf:NetworkInterface{id: elb_association.netinf_id}),
-            (elb:LoadBalancerV2{id: elb_association.elb_id})
+            (elb:LoadBalancerV2{id: elb_association.elb_arn})
         MERGE (elb)-[r:NETWORK_INTERFACE]->(netinf)
         ON CREATE SET r.firstseen = timestamp()
         SET r.lastupdated = $update_tag
@@ -235,11 +235,17 @@ def load(neo4j_session: neo4j.Session, data: List[Dict], region: str, aws_accoun
 
     for network_interface in data:
         # https://aws.amazon.com/premiumsupport/knowledge-center/elb-find-load-balancer-IP/
-        matchObj = re.match(r'^ELB (?:net|app)/([^\/]+)\/(.*)', network_interface.get('Description', ''))
+        matchObj = re.match(r'^ELB (?:net|app|gwy)/([^\/]+)\/(.*)', network_interface.get('Description', ''))
         if matchObj:
+            # get the end of arn from network interface description
+            elb_name_id: str = network_interface.get('Description', '').split(' ')[1]
+            # ELBV2 arn that is id of every LoadBalancerV2 and will be used to make
+            # (:LoadBalancerV2)-[:NETWORK_INTERFACE]->(:NetworkInterface)
+            elb_arn = f'arn:aws:elasticloadbalancing:{region}:{aws_account_id}:loadbalancer/{elb_name_id}'
             elb_associations_v2.append({
                 'netinf_id': network_interface['NetworkInterfaceId'],
-                'elb_id': f'{matchObj[1]}-{matchObj[2]}.elb.{region}.amazonaws.com',
+                'elb_dnsname': f'{matchObj[1]}-{matchObj[2]}.elb.{region}.amazonaws.com',
+                'elb_arn': elb_arn,
             })
         else:
             matchObj = re.match(r'^ELB (.*)', network_interface.get('Description', ''))

tests/data/aws/ec2/load_balancers.py

@@ -19,6 +19,12 @@
         'Protocol': 'HTTPS',
         'TargetGroupArn': 'arn:aws:ec2:us-east-1:012345678912:targetgroup',
     },
+    {
+        'ListenerArn': "arn:aws:elasticloadbalancing:us-east-1:000000000000:listener/gwy/mytestgwy/gwyLBId/gwyListId",
+        'Port': 500,
+        'Protocol': 'GENEVE',
+        'TargetGroupArn': 'arn:aws:ec2:us-east-1:012345678912:targetgroup',
+    },
 ]
 
 # Listener fields
@@ -89,6 +95,10 @@
 LOAD_BALANCER_DATA = [
     {
         'DNSName': 'myawesomeloadbalancer.amazonaws.com',
+        'LoadBalancerArn': (
+            "arn:aws:ec2:elasticloadbalancing:us-east-1:000000000000:"
+            "loadbalancer/app/myawesomeloadbalancer/someid"
+        ),
         'CreatedTime': '10-27-2019 12:35AM',
         'LoadBalancerName': 'myawesomeloadbalancer',
         'Type': 'application',
@@ -109,6 +119,34 @@
         'Listeners': LOAD_BALANCER_LISTENERS,
         'TargetGroups': TARGET_GROUPS,
     },
+    {
+        'LoadBalancerArn': (
+            'arn:aws:elasticloadbalancing:eu-north-1:167992319538'
+            ':loadbalancer/gwy/test-gateway-load-balancer/180ff0c1e66f6754'
+        ),
+        'CreatedTime': '2023-07-14 17:27:50.495000+00:00',
+        'LoadBalancerName': 'test-gateway-load-balancer',
+        'VpcId': 'vpc-03e880ef713e1f725',
+        'State': {
+            'Code': 'active',
+        },
+        'Type': 'gateway',
+        'AvailabilityZones': [
+            {
+                'ZoneName': 'myAZ',
+                'SubnetId': 'mysubnetIdA',
+                'LoadBalancerAddresses': [
+                    {
+                        'IpAddress': '50.0.1.0',
+                        'AllocationId': 'someId',
+                    },
+                ],
+            },
+        ],
+        'IpAddressType': 'ipv4',
+        'Listeners': LOAD_BALANCER_LISTENERS,
+        'TargetGroups': TARGET_GROUPS,
+    },
 ]
 
 # 'LoadBalancers': [

tests/integration/cartography/intel/aws/ec2/test_ec2_load_balancers.py

@@ -6,12 +6,84 @@
 TEST_UPDATE_TAG = 123456789
 
 
+def test_load_gwy_load_balancers_v2(neo4j_session, *args):
+    # gateway load balancers do not have security groups
+    load_balancer_data = tests.data.aws.ec2.load_balancers.LOAD_BALANCER_DATA
+    ec2_instance_id = 'i-0f76fade'
+    gwy_load_balancer_id = (
+        'arn:aws:elasticloadbalancing:eu-north-1:167992319538:'
+        'loadbalancer/gwy/test-gateway-load-balancer/180ff0c1e66f6754'
+    )
+
+    # an ec2instance and AWSAccount must exist but Security Groups aren't needed since
+    # gateway type load balancers do not have them
+    neo4j_session.run(
+        """
+        MERGE (ec2:EC2Instance{instanceid: $ec2_instance_id})
+        ON CREATE SET ec2.firstseen = timestamp()
+        SET ec2.lastupdated = $aws_update_tag
+
+        MERGE (aws:AWSAccount{id: $aws_account_id})
+        ON CREATE SET aws.firstseen = timestamp()
+        SET aws.lastupdated = $aws_update_tag
+        """,
+        ec2_instance_id=ec2_instance_id,
+        aws_account_id=TEST_ACCOUNT_ID,
+        aws_update_tag=TEST_UPDATE_TAG,
+    )
+
+    # Makes elbv2
+    # (aa)-[r:RESOURCE]->(elbv2)
+    # also makes
+    # (elbv2)->[RESOURCE]->(EC2Subnet)
+    # should not make relationship with SecurityGroup
+    cartography.intel.aws.ec2.load_balancer_v2s.load_load_balancer_v2s(
+        neo4j_session,
+        load_balancer_data,
+        TEST_REGION,
+        TEST_ACCOUNT_ID,
+        TEST_UPDATE_TAG,
+    )
+
+    nodes = neo4j_session.run(
+        """
+        MATCH (aa:AWSAccount{id: $AWS_ACCOUNT_ID})
+            -[r1:RESOURCE]->(elbv2:LoadBalancerV2{id: $ID})
+            -[r2:ELBV2_LISTENER]->(l:ELBV2Listener{id: $LISTENER_ARN})
+        RETURN aa.id, elbv2.id, l.id
+        """,
+        AWS_ACCOUNT_ID=TEST_ACCOUNT_ID,
+        ID=gwy_load_balancer_id,
+        LISTENER_ARN="arn:aws:elasticloadbalancing:us-east-1:000000000000:listener/gwy/mytestgwy/gwyLBId/gwyListId",
+    )
+
+    expected_nodes = {
+        (
+            TEST_ACCOUNT_ID,
+            gwy_load_balancer_id,
+            "arn:aws:elasticloadbalancing:us-east-1:000000000000:listener/gwy/mytestgwy/gwyLBId/gwyListId",
+        ),
+    }
+    actual_nodes = {
+        (
+            n['aa.id'],
+            n['elbv2.id'],
+            n['l.id'],
+        )
+        for n in nodes
+    }
+    assert actual_nodes == expected_nodes
+
+
 def test_load_load_balancer_v2s(neo4j_session, *args):
     load_balancer_data = tests.data.aws.ec2.load_balancers.LOAD_BALANCER_DATA
     ec2_instance_id = 'i-0f76fade'
     sg_group_id = 'sg-123456'
     sg_group_id_2 = 'sg-234567'
-    load_balancer_id = "myawesomeloadbalancer.amazonaws.com"
+    load_balancer_id = (
+        "arn:aws:ec2:elasticloadbalancing:us-east-1:000000000000:"
+        "loadbalancer/app/myawesomeloadbalancer/someid"
+    )
 
     # an ec2instance and AWSAccount must exist
     neo4j_session.run(

tests/integration/cartography/intel/aws/test_route53.py

@@ -102,7 +102,7 @@ def test_load_dnspointsto_ec2_relationships(neo4j_session):
     result = neo4j_session.run(
         """
         MATCH (n:AWSDNSRecord{id:"/hostedzone/HOSTED_ZONE/elbv2.example.com/ALIAS"})
-        -[:DNS_POINTS_TO]->(l:LoadBalancerV2{id:"myawesomeloadbalancer.amazonaws.com"})
+        -[:DNS_POINTS_TO]->(l:LoadBalancerV2{dnsname:"myawesomeloadbalancer.amazonaws.com"})
         return n.name, l.name
         """,
     )