Merge pull request #320 from nccgroup/enhancement/additional-error-lo…

…gging

Enhancement/Additional AWS exception handling take 2

ScoutSuite/providers/aws/facade/emr.py

@@ -8,13 +8,13 @@ class EMRFacade(AWSBaseFacade):
     async def get_clusters(self, region):
         cluster_list = await AWSFacadeUtils.get_all_pages('emr', region, self.session, 'list_clusters', 'Clusters')
         cluster_ids = [cluster['Id'] for cluster in cluster_list]
-        client = AWSFacadeUtils.get_client('emr', self.session, region)
 
+        return await map_concurrently(self._get_cluster, cluster_ids, region=region)
+
+    async def _get_cluster(self, cluster_id: str, region: str):
+        client = AWSFacadeUtils.get_client('emr', self.session, region)
         try:
-            return await map_concurrently(
-                lambda cluster_id: run_concurrently(
-                    lambda: client.describe_cluster(ClusterId=cluster_id)['Cluster']),
-                cluster_ids)
+            return await run_concurrently(lambda: client.describe_cluster(ClusterId=cluster_id)['Cluster'])
         except Exception as e:
             print_exception('Failed to describe EMR cluster: {}'.format(e))
-            return []
+            raise

ScoutSuite/providers/aws/facade/ses.py

@@ -21,7 +21,7 @@ async def _get_identity_dkim_attributes(self, identity_name: str, region: str):
             )
         except Exception as e:
             print_exception('Failed to get SES DKIM attributes: {}'.format(e))
-            dkim_attributes = None
+            raise
         return identity_name, dkim_attributes
 
     async def get_identity_policies(self, region: str, identity_name: str):

ScoutSuite/providers/aws/facade/sqs.py

@@ -29,6 +29,6 @@ async def _get_queue_attributes(self, queue_url: str, region: str, attribute_nam
             )
         except Exception as e:
             print_exception('Failed to get SQS queue attributes: {}'.format(e))
-            queue_attributes = None
+            raise
 
         return queue_url, queue_attributes

ScoutSuite/providers/aws/resources/directconnect/base.py

@@ -6,10 +6,9 @@
 class Connections(AWSResources):
     async def fetch_all(self, **kwargs):
         raw_connections = await self.facade.directconnect.get_connections(self.scope['region'])
-        if raw_connections:
-            for raw_connection in raw_connections:
-                name, resource = self._parse_function(raw_connection)
-                self[name] = resource
+        for raw_connection in raw_connections:
+            name, resource = self._parse_function(raw_connection)
+            self[name] = resource
 
     def _parse_function(self, raw_connection):
         raw_connection['id'] = raw_connection.pop('connectionId')

ScoutSuite/providers/aws/resources/ec2/ami.py

@@ -4,10 +4,9 @@
 class AmazonMachineImages(AWSResources):
     async def fetch_all(self, **kwargs):
         raw_images = await self.facade.ec2.get_images(self.scope['region'], self.scope['owner_id'])
-        if raw_images:
-            for raw_image in raw_images:
-                name, resource = self._parse_image(raw_image)
-                self[name] = resource
+        for raw_image in raw_images:
+            name, resource = self._parse_image(raw_image)
+            self[name] = resource
 
     def _parse_image(self, raw_image):
         raw_image['id'] = raw_image['ImageId']

ScoutSuite/providers/aws/resources/vpcs.py

@@ -13,25 +13,17 @@ def __init__(self, facade, scope: dict, add_ec2_classic=False):
         self.add_ec2_classic = add_ec2_classic
 
     async def fetch_all(self, **kwargs):
+        raw_vpcs = await self.facade.ec2.get_vpcs(self.scope['region'])
 
-        try:
-            raw_vpcs = await self.facade.ec2.get_vpcs(self.scope['region'])
-        except Exception as e:
-            print_exception('Failed to get VPCs for region {}: {}'.format(self.scope['region'], e))
-        else:
-            if raw_vpcs:
-                for raw_vpc in raw_vpcs:
-                    vpc_id, vpc = self._parse_vpc(raw_vpc)
-                    self[vpc_id] = vpc
+        for raw_vpc in raw_vpcs:
+            vpc_id, vpc = self._parse_vpc(raw_vpc)
+            self[vpc_id] = vpc
 
-                try:
-                    await self._fetch_children_of_all_resources(
-                        resources=self,
-                        scopes={vpc_id: {'region': self.scope['region'], 'vpc': vpc_id}
-                                for vpc_id in self}
-                    )
-                except Exception as e:
-                    print_exception('Failed to fetch resources for VPC {}: {}'.format(vpc_id, e))
+        await self._fetch_children_of_all_resources(
+            resources=self,
+            scopes={vpc_id: {'region': self.scope['region'], 'vpc': vpc_id}
+                    for vpc_id in self}
+        )
 
     def _parse_vpc(self, vpc):
         return vpc['VpcId'], {}

ScoutSuite/providers/utils.py

@@ -1,8 +1,6 @@
 import asyncio
 from hashlib import sha1
 
-from ScoutSuite.core.console import print_exception
-
 
 def get_non_provider_id(name):
     """
@@ -17,16 +15,16 @@ def get_non_provider_id(name):
     return name_hash.hexdigest()
 
 
-def run_concurrently(func):
+def run_concurrently(function):
     """
-    Schedules the execution of function `func` in the default thread pool (referred as 'executor') that has been
+    Schedules the execution of function `function` in the default thread pool (referred as 'executor') that has been
     associated with the global event loop.
 
-    :param func: function to be executed concurrently, in a dedicated thread.
+    :param function: function to be executed concurrently, in a dedicated thread.
     :return: an asyncio.Future to be awaited.
     """
 
-    return asyncio.get_event_loop().run_in_executor(executor=None, func=func)
+    return asyncio.get_event_loop().run_in_executor(executor=None, func=function)
 
 
 async def get_and_set_concurrently(get_and_set_funcs: [], entities: [], **kwargs):
@@ -38,31 +36,32 @@ async def get_and_set_concurrently(get_and_set_funcs: [], entities: [], **kwargs
     :param get_and_set_funcs: list of functions that takes a region and an entity (they must have the following
     signature: region: str, entity: {}) and then fetch and set some kind of attributes to this entity.
     :param entities: list of a same kind of entities
-    :param region: a region
+    :param kwargs: used to pass cloud provider specific parameters (ex: region or vpc for AWS, etc.) to the given
+    functions.
 
     :return:
     """
 
     if len(entities) == 0:
         return
 
-    try:
-        tasks = {
-            asyncio.ensure_future(
-                get_and_set_func(entity, **kwargs)
-            ) for entity in entities for get_and_set_func in get_and_set_funcs
-        }
-        await asyncio.wait(tasks)
-    except Exception as e:
-        print_exception('Failed to run function concurrently: {}'.format(e))
+    tasks = {
+        asyncio.ensure_future(
+            get_and_set_func(entity, **kwargs)
+        ) for entity in entities for get_and_set_func in get_and_set_funcs
+    }
+    await asyncio.wait(tasks)
 
 
-async def map_concurrently(coro, entities, **kwargs):
+async def map_concurrently(coroutine, entities, **kwargs):
     """
-    Given a list of entities, executes coroutine `coro` concurrently on each entity and returns a list of the obtained
-    results ([await coro(entity_x), await coro(entity_a), ..., await coro(entity_z)]).
+    Given a list of entities, executes coroutine `coroutine` concurrently on each entity and returns a list of the
+    obtained results ([await coroutine(entity_x), await coroutine(entity_a), ..., await coroutine(entity_z)]).
 
-    :param coro: coroutine to be executed concurrently. Takes an entity as parameter and returns a new entity.
+    :param coroutine: coroutine to be executed concurrently. Takes an entity as parameter and returns a new entity.
+    If the given coroutine does some exception handling, it should ensure to propagate the handled exceptions so
+    `map_concurrently` can handle them as well (in particular ignoring them) to avoid `None` values in the list
+    returned.
     :param entities: a list of the same type of entity (ex: cluster ids)
 
     :return: a list of new entities (ex: clusters)
@@ -72,16 +71,19 @@ async def map_concurrently(coro, entities, **kwargs):
         return []
 
     results = []
-    try:
-        tasks = {
-            asyncio.ensure_future(
-                coro(entity, **kwargs)
-            ) for entity in entities
-        }
-        for task in asyncio.as_completed(tasks):
+
+    tasks = {
+        asyncio.ensure_future(
+            coroutine(entity, **kwargs)
+        ) for entity in entities
+    }
+
+    for task in asyncio.as_completed(tasks):
+        try:
             result = await task
+        except Exception:
+            pass
+        else:
             results.append(result)
-    except Exception as e:
-        print_exception('Failed to map concurrently: {}'.format(e))
 
     return results