From 0b34d30fd6e7908df27de705b9272eb109e7f7ea Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Tue, 6 Apr 2021 11:35:41 +0300 Subject: [PATCH] Improved exploiter descriptors and usage of them in report.py --- .../exploit_processing/exploiter_descriptor_enum.py | 13 +++---------- .../monkey_island/cc/services/reporting/report.py | 9 ++++----- 2 files changed, 7 insertions(+), 15 deletions(-) diff --git a/monkey/monkey_island/cc/services/reporting/issue_processing/exploit_processing/exploiter_descriptor_enum.py b/monkey/monkey_island/cc/services/reporting/issue_processing/exploit_processing/exploiter_descriptor_enum.py index 08c471174..65964b5de 100644 --- a/monkey/monkey_island/cc/services/reporting/issue_processing/exploit_processing/exploiter_descriptor_enum.py +++ b/monkey/monkey_island/cc/services/reporting/issue_processing/exploit_processing/exploiter_descriptor_enum.py @@ -36,14 +36,7 @@ class ExploiterDescriptorEnum(Enum): ZEROLOGON = ExploiterDescriptor('ZerologonExploiter', 'Zerologon Exploiter', ZerologonExploitProcessor) @staticmethod - def __dict__() -> Dict[str, ExploiterDescriptor]: - descriptor_dict = {} - for descriptor in ExploiterDescriptorEnum: - descriptor_dict[descriptor.value.class_name] = descriptor - return descriptor_dict - - @staticmethod - def get_display_name_by_class_name(class_name: str) -> str: - return [descriptor.display_name.value + def get_by_class_name(class_name: str) -> ExploiterDescriptor: + return [descriptor.value for descriptor in ExploiterDescriptorEnum - if descriptor.class_name.value == class_name][0] + if descriptor.value.class_name == class_name][0] diff --git a/monkey/monkey_island/cc/services/reporting/report.py b/monkey/monkey_island/cc/services/reporting/report.py index d84c5ad94..6430a2559 100644 --- a/monkey/monkey_island/cc/services/reporting/report.py +++ b/monkey/monkey_island/cc/services/reporting/report.py @@ -32,7 +32,6 @@ logger = logging.getLogger(__name__) class ReportService: - exploiter_descriptors = ExploiterDescriptorEnum.__dict__ class DerivedIssueEnum: WEAK_PASSWORD = "weak_password" @@ -145,7 +144,7 @@ class ReportService: @staticmethod def get_exploits_used_on_node(node: dict) -> List[str]: - return list(set([ExploiterDescriptorEnum.get_display_name_by_class_name(exploit['exploiter']) + return list(set([ExploiterDescriptorEnum.get_by_class_name(exploit['exploiter']).display_name for exploit in node['exploits'] if exploit['result']])) @@ -253,9 +252,9 @@ class ReportService: @staticmethod def process_exploit(exploit) -> ExploiterReportInfo: exploiter_type = exploit['data']['exploiter'] - exploiter_descriptor = ReportService.exploiter_descriptors[exploiter_type].value + exploiter_descriptor = ExploiterDescriptorEnum.get_by_class_name(exploiter_type) processor = exploiter_descriptor.processor() - exploiter_info = processor.get_exploit_info_by_dict(exploiter_descriptor.class_name, exploit) + exploiter_info = processor.get_exploit_info_by_dict(exploiter_type, exploit) return exploiter_info @staticmethod @@ -485,7 +484,7 @@ class ReportService: if exploits == default_exploits: return ['default'] - return [ExploiterDescriptorEnum.get_display_name_by_class_name(exploit) + return [ExploiterDescriptorEnum.get_by_class_name(exploit).display_name for exploit in exploits] @staticmethod