Fix KeyError exception when getting DNSName for aws elbv2 (#996)

lyft · Jul 14, 2023 · c02da83 · c02da83
1 parent 48f50ca
commit c02da83
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/cartography/intel/aws/ec2/load_balancer_v2s.py b/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"),

diff --git a/cartography/intel/aws/ec2/network_interfaces.py b/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 = 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', ''))