From d2181f6577ee34d9cd39c984d6895f86641eed69 Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Wed, 5 Jan 2022 14:15:56 +0200
Subject: [PATCH] Agent, UT: fix ldap builder UT's and some imports in
 log4shell

---
 monkey/infection_monkey/exploit/log4shell.py             | 9 ++++++---
 .../infection_monkey/exploit/log4shell_utils/__init__.py | 7 ++++++-
 .../exploit/log4shell_utils/test_exploit_builder.py      | 3 ++-
 3 files changed, 14 insertions(+), 5 deletions(-)

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