diff --git a/monkey/infection_monkey/exploit/log4shell.py b/monkey/infection_monkey/exploit/log4shell.py index a8af38401..d3dc9d351 100644 --- a/monkey/infection_monkey/exploit/log4shell.py +++ b/monkey/infection_monkey/exploit/log4shell.py @@ -4,10 +4,11 @@ import logging from threading import Thread from common.utils.exploit_enum import ExploitType -from infection_monkey.exploit.log4shell_utils import LDAPExploitServer, build_exploit_bytecode -from infection_monkey.exploit.log4shell_utils.exploit_builder import ( +from infection_monkey.exploit.log4shell_utils import ( LINUX_EXPLOIT_TEMPLATE_PATH, WINDOWS_EXPLOIT_TEMPLATE_PATH, + LDAPExploitServer, + build_exploit_bytecode, ) from infection_monkey.exploit.log4shell_utils.requests import exploits from infection_monkey.exploit.tools.helpers import get_monkey_depth @@ -133,7 +134,9 @@ class Log4ShellExploiter(WebRCE): Log4ShellExploiter.HTTPHandler.class_downloaded = True def _run_class_http_server(self, ip): - server = http.server.HTTPServer((ip, self.class_http_server_port), Log4ShellExploiter.HTTPHandler) + server = http.server.HTTPServer( + (ip, self.class_http_server_port), Log4ShellExploiter.HTTPHandler + ) while ( not Log4ShellExploiter.HTTPHandler.class_downloaded and not Log4ShellExploiter.HTTPHandler.stop diff --git a/monkey/infection_monkey/exploit/log4shell_utils/__init__.py b/monkey/infection_monkey/exploit/log4shell_utils/__init__.py index 5ffe79e49..d57218d8d 100644 --- a/monkey/infection_monkey/exploit/log4shell_utils/__init__.py +++ b/monkey/infection_monkey/exploit/log4shell_utils/__init__.py @@ -1,2 +1,7 @@ -from .exploit_builder import build_exploit_bytecode, InvalidExploitTemplateError +from .exploit_builder import ( + build_exploit_bytecode, + InvalidExploitTemplateError, + LINUX_EXPLOIT_TEMPLATE_PATH, + WINDOWS_EXPLOIT_TEMPLATE_PATH, +) from .ldap_server import LDAPExploitServer diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/log4shell_utils/test_exploit_builder.py b/monkey/tests/unit_tests/infection_monkey/exploit/log4shell_utils/test_exploit_builder.py index e2551d061..01d4f61c4 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/log4shell_utils/test_exploit_builder.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/log4shell_utils/test_exploit_builder.py @@ -1,6 +1,7 @@ import pytest from infection_monkey.exploit.log4shell_utils import ( + LINUX_EXPLOIT_TEMPLATE_PATH, InvalidExploitTemplateError, build_exploit_bytecode, ) @@ -19,7 +20,7 @@ def invalid_templates_dir(data_for_tests_dir): def test_inject_command(payload): expected_bytecode = b"\x21" + payload.encode() + b"\x0c" - exploit_bytecode = build_exploit_bytecode(payload) + exploit_bytecode = build_exploit_bytecode(payload, LINUX_EXPLOIT_TEMPLATE_PATH) assert expected_bytecode in exploit_bytecode