Agent, UT: fix ldap builder UT's and some imports in log4shell

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

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

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