Minor refactorings of code style in zero trust code

monkey/common/utils/exceptions.py

@@ -44,3 +44,7 @@ class UnknownFindingError(Exception):
 
 class VersionServerConnectionError(Exception):
     """ Raise to indicate that connection to version update server failed """
+
+
+class FindingWithoutDetailsError(Exception):
+    """ Raise when pulling events for a finding, but get none """

monkey/infection_monkey/exploit/HostExploiter.py

@@ -74,10 +74,10 @@ class HostExploiter(Plugin):
         result = None
         try:
             result = self._exploit_host()
-        except FailedExploitationError as e:
-            logger.debug(e)
-        except Exception as _:
-            logger.error(f'Exception in exploit_host', exc_info=True)
+        except FailedExploitationError:
+            logger.debug('Exploiter failed.', exc_info=True)
+        except Exception:
+            logger.error('Exception in exploit_host', exc_info=True)
         finally:
             self.post_exploit()
         return result

monkey/infection_monkey/system_info/collectors/scoutsuite_collector/__init__.py

@@ -0,0 +1,15 @@
+import pkgutil
+import sys
+from pathlib import PurePath
+
+_scoutsuite_api_package = pkgutil.get_loader('common.cloud.scoutsuite.ScoutSuite.__main__')
+
+
+def _add_scoutsuite_to_python_path():
+    scoutsuite_path = PurePath(_scoutsuite_api_package.path).parent.parent.__str__()
+    sys.path.append(scoutsuite_path)
+
+
+# Add ScoutSuite to python path because this way
+# we don't need to change any imports in ScoutSuite code
+_add_scoutsuite_to_python_path()

monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_api.py

@@ -1,19 +1,3 @@
-import pkgutil
-import sys
-from pathlib import PurePath
-
-_scoutsuite_api_package = pkgutil.get_loader('common.cloud.scoutsuite.ScoutSuite.__main__')
-
-
-def _add_scoutsuite_to_python_path():
-    scoutsuite_path = PurePath(_scoutsuite_api_package.path).parent.parent.__str__()
-    sys.path.append(scoutsuite_path)
-
-
-# Add ScoutSuite to python path because this way
-# we don't need to change any imports in ScoutSuite code
-_add_scoutsuite_to_python_path()
-
 import common.cloud.scoutsuite.ScoutSuite.api_run as scoutsuite_api
 
 

monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_collector.py

@@ -1,7 +1,8 @@
 import logging
+from typing import Union
 
 import infection_monkey.system_info.collectors.scoutsuite_collector.scoutsuite_api as scoutsuite_api
-from common.cloud.scoutsuite.ScoutSuite.providers.aws.provider import AWSProvider
+from common.cloud.scoutsuite.ScoutSuite.providers.base.provider import BaseProvider
 from common.cloud.scoutsuite_consts import CloudProviders
 from common.utils.exceptions import ScoutSuiteScanError
 from infection_monkey.config import WormConfiguration
@@ -20,12 +21,12 @@ def scan_cloud_security(cloud_type: CloudProviders):
         logger.error(f"ScoutSuite didn't scan {cloud_type.value} security because: {e}")
 
 
-def run_scoutsuite(cloud_type: str):
+def run_scoutsuite(cloud_type: str) -> Union[BaseProvider, dict]:
     return scoutsuite_api.run(provider=cloud_type,
                               aws_access_key_id=WormConfiguration.aws_access_key_id,
                               aws_secret_access_key=WormConfiguration.aws_secret_access_key,
                               aws_session_token=WormConfiguration.aws_session_token)
 
 
-def send_results(results: AWSProvider):
+def send_results(results: BaseProvider):
     ScoutSuiteTelem(results).send()

monkey/infection_monkey/telemetry/scoutsuite_telem.py

@@ -1,12 +1,12 @@
 from common.cloud.scoutsuite.ScoutSuite.output.result_encoder import ScoutJsonEncoder
-from common.cloud.scoutsuite.ScoutSuite.providers.aws.provider import AWSProvider
+from common.cloud.scoutsuite.ScoutSuite.providers.base.provider import BaseProvider
 from common.common_consts.telem_categories import TelemCategoryEnum
 from infection_monkey.telemetry.base_telem import BaseTelem
 
 
 class ScoutSuiteTelem(BaseTelem):
 
-    def __init__(self, data: AWSProvider):
+    def __init__(self, data: BaseProvider):
         """
         Default ScoutSuite telemetry constructor
         :param data: Data gathered via ScoutSuite

monkey/monkey_island/cc/services/zero_trust/monkey_findings/monkey_zt_details_service.py

@@ -2,6 +2,7 @@ from typing import List
 
 from bson import ObjectId
 
+from common.utils.exceptions import FindingWithoutDetailsError
 from monkey_island.cc.models.zero_trust.monkey_finding_details import MonkeyFindingDetails
 
 
@@ -19,14 +20,13 @@ class MonkeyZTDetailsService:
                                     'latest_events': {'$slice': ['$events', int(-1 * MAX_EVENT_FETCH_CNT / 2)]},
                                     'event_count': {'$size': '$events'}}},
                     {'$unset': ['events']}]
-        details = list(MonkeyFindingDetails.objects.aggregate(*pipeline))
+        details = list(MonkeyFindingDetails.objects.aggregate(*pipeline))[0]
         if details:
-            details = details[0]
             details['latest_events'] = MonkeyZTDetailsService._remove_redundant_events(details['event_count'],
                                                                                        details['latest_events'])
             return details
         else:
-            return {}
+            raise FindingWithoutDetailsError(f"Finding {finding_id} had no details.")
 
     @staticmethod
     def _remove_redundant_events(fetched_event_count: int, latest_events: List[object]) -> List[object]: