Merge pull request duo-labs#104 from 0xdabbad00/fix_elbv2

Fix ELBv2s in prepare

Author: 0xdabbad00

commands/prepare.py

@@ -83,7 +83,7 @@ def get_elbs(subnet):
 
     # ALBs and NLBs
     alb_instances = query_aws(subnet.account, "elbv2-describe-load-balancers", subnet.region)
-    alb_resource_filter = '.LoadBalancers[] | select(.VPCId == "{}") | select(.Subnets[] == "{}")'
+    alb_resource_filter = '.LoadBalancers[] | select(.VpcId == "{}") | select(.AvailabilityZones[].SubnetId == "{}")'
     albs = pyjq.all(alb_resource_filter.format(subnet.vpc.local_id, subnet.local_id), alb_instances)
 
     return elbs + albs
@@ -250,7 +250,6 @@ def build_data_structure(account_data, config, outputfilter):
                         elb = Elb(subnet, elb_json)
                         subnet.addChild(elb)
 
-
                     # If there are leaves, then add this subnet to the final graph
                     if len(subnet.leaves) > 0:
                         node_count_per_region += len(subnet.leaves)

tests/unit/test_prepare.py

@@ -63,13 +63,13 @@ def test_build_data_structure(self):
 
         # Now check it
         # Check number of connections
-        assert_equal(14, len(pyjq.all('.[].data|select(.type == "edge")|keys', cytoscape_json)))
+        assert_equal(17, len(pyjq.all('.[].data|select(.type == "edge")|keys', cytoscape_json)))
 
         # Check number of nodes
         assert_equal(2, len(pyjq.all('.[].data|select(.type == "ip")|keys', cytoscape_json)))
         assert_equal(2, len(pyjq.all('.[].data|select(.type == "rds")|keys', cytoscape_json)))
         assert_equal(3, len(pyjq.all('.[].data|select(.type == "ec2")|keys', cytoscape_json)))
-        assert_equal(2, len(pyjq.all('.[].data|select(.type == "elb")|keys', cytoscape_json)))
+        assert_equal(3, len(pyjq.all('.[].data|select(.type == "elb")|keys', cytoscape_json)))
         assert_equal(4, len(pyjq.all('.[].data|select(.type == "subnet")|keys', cytoscape_json)))
         assert_equal(1, len(pyjq.all('.[].data|select(.type == "region")|keys', cytoscape_json)))
         assert_equal(1, len(pyjq.all('.[].data|select(.type == "vpc")|keys', cytoscape_json)))