From 6f16ba431cecc054695c8e2572fa15985484d548 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Tue, 9 Feb 2021 14:21:55 +0200 Subject: [PATCH] Minor refactorings of code style in zero trust code --- monkey/common/utils/exceptions.py | 4 ++++ monkey/infection_monkey/exploit/HostExploiter.py | 8 ++++---- .../collectors/scoutsuite_collector/__init__.py | 15 +++++++++++++++ .../scoutsuite_collector/scoutsuite_api.py | 16 ---------------- .../scoutsuite_collector/scoutsuite_collector.py | 7 ++++--- .../telemetry/scoutsuite_telem.py | 4 ++-- .../monkey_findings/monkey_zt_details_service.py | 6 +++--- 7 files changed, 32 insertions(+), 28 deletions(-) diff --git a/monkey/common/utils/exceptions.py b/monkey/common/utils/exceptions.py index 6b992b557..2c7121942 100644 --- a/monkey/common/utils/exceptions.py +++ b/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 """ diff --git a/monkey/infection_monkey/exploit/HostExploiter.py b/monkey/infection_monkey/exploit/HostExploiter.py index f8fda6d18..c48cadcf0 100644 --- a/monkey/infection_monkey/exploit/HostExploiter.py +++ b/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 diff --git a/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/__init__.py b/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/__init__.py index e69de29bb..97e736b4b 100644 --- a/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/__init__.py +++ b/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() diff --git a/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_api.py b/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_api.py index 575058946..88ef32293 100644 --- a/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_api.py +++ b/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 diff --git a/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_collector.py b/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_collector.py index 7726e980e..0664b9b0f 100644 --- a/monkey/infection_monkey/system_info/collectors/scoutsuite_collector/scoutsuite_collector.py +++ b/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() diff --git a/monkey/infection_monkey/telemetry/scoutsuite_telem.py b/monkey/infection_monkey/telemetry/scoutsuite_telem.py index 4e49c0695..5ad553e94 100644 --- a/monkey/infection_monkey/telemetry/scoutsuite_telem.py +++ b/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 diff --git a/monkey/monkey_island/cc/services/zero_trust/monkey_findings/monkey_zt_details_service.py b/monkey/monkey_island/cc/services/zero_trust/monkey_findings/monkey_zt_details_service.py index 63f809fee..3b2a0eed0 100644 --- a/monkey/monkey_island/cc/services/zero_trust/monkey_findings/monkey_zt_details_service.py +++ b/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]: