Add ELF loader

dissect/executable/elf/tools/loader.py

@@ -0,0 +1,380 @@
+from __future__ import annotations
+
+import argparse
+import ctypes
+import ctypes.util
+import enum
+import json
+import logging
+import os
+import platform
+import random
+import struct
+import sys
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, BinaryIO
+
+from dissect.executable.elf.c_elf import PF, PT
+from dissect.executable.elf.elf import ELF
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from dissect.executable.elf.elf import Segment
+
+log = logging.getLogger(__name__)
+log.setLevel(os.getenv("DISSECT_LOG_EXECUTABLE", "CRITICAL"))
+
+
+class PROT(enum.IntFlag):
+    READ = 0x1
+    WRITE = 0x2
+    EXEC = 0x4
+
+    @classmethod
+    def from_PF(cls, flags: PF) -> PROT:
+        mapping = {
+            PF.R: PROT.READ,
+            PF.X: PROT.EXEC,
+            PF.W: PROT.WRITE,
+        }
+
+        result = 0
+        for pf, prot in mapping.items():
+            if flags & pf:
+                result |= prot
+
+        return result
+
+
+class MAP(enum.IntFlag):
+    SHARED = 0x01
+    PRIVATE = 0x02
+    FIXED = 0x10
+    ANONYMOUS = 0x20
+
+
+class AT(enum.IntEnum):
+    NULL = 0
+    IGNORE = 1
+    EXECFD = 2
+    PHDR = 3
+    PHENT = 4
+    PHNUM = 5
+    PAGESZ = 6
+    BASE = 7
+    FLAGS = 8
+    ENTRY = 9
+    NOTELF = 10
+    UID = 11
+    EUID = 12
+    GID = 13
+    EGID = 14
+    PLATFORM = 15
+    HWCAP = 16
+    CLKTCK = 17
+    SECURE = 23
+    BASE_PLATFORM = 24
+    RANDOM = 25
+    EXECFN = 31
+    SYSINFO = 32
+    SYSINFO_EHDR = 33
+
+
+PAGE_SIZE = 4096
+ALIGN = PAGE_SIZE - 1
+
+
+libc = None
+loader = None
+
+
+def _setup_libc(path: str | None = None) -> None:
+    """Set up the libc functions.
+
+    Args:
+        path: The optional path to the libc library to use.
+    """
+    path = path or ctypes.util.find_library("c")
+    if path is None:
+        raise ValueError("Unable to find path to libc")
+
+    global libc
+    if libc is not None:
+        return
+
+    libc = ctypes.CDLL(path, use_errno=True)
+
+    libc.mmap.argtypes = [
+        # addr
+        ctypes.c_size_t,
+        # length
+        ctypes.c_size_t,
+        # prot
+        ctypes.c_int,
+        # flags
+        ctypes.c_int,
+        # fd
+        ctypes.c_int,
+        # offset
+        ctypes.c_size_t,
+    ]
+    libc.mmap.restype = ctypes.c_size_t
+
+    libc.mprotect.argtypes = [
+        # addr
+        ctypes.c_size_t,
+        # len
+        ctypes.c_size_t,
+        # prot
+        ctypes.c_int,
+    ]
+    libc.mprotect.restype = ctypes.c_int
+
+    libc.getauxval.argtypes = [
+        # type
+        ctypes.c_ulong
+    ]
+    libc.getauxval.restype = ctypes.c_size_t
+
+    if log.level <= logging.DEBUG:
+        libc.mmap = _ctype_log(libc.mmap)
+        libc.mprotect = _ctype_log(libc.mprotect)
+        libc.getauxval = _ctype_log(libc.getauxval)
+
+
+def _setup_loader() -> None:
+    """Set up the loader function."""
+    global loader
+    if loader is not None:
+        return
+
+    machine = platform.machine()
+    if machine == "x86_64":
+        asm = bytes.fromhex(
+            # rdi = pointer to initial stack
+            # rsi = size of stack
+            # rdx = entry point
+            #
+            # Reserve space for the stack and align it
+            "   48 29 f4"  # sub rsp, rsi
+            "48 83 e4 f0"  # and rsp, 0fffffffffffffff0h
+            # Copy the stack over (dst=rdi, src=rsi, count=rcx)
+            "   48 89 f1"  # mov rcx, rsi
+            "   48 89 fe"  # mov rsi, rdi
+            "   48 89 e7"  # mov rdi, rsp
+            "         fc"  # cld
+            "      f3 a4"  # rep movsb
+            # Jump to the entry point
+            "      ff e2"  # jmp rdx
+            "         f4"  # hlt
+        )
+    elif machine == "aarch64":
+        asm = bytes.fromhex(
+            # x0 = pointer to initial stack
+            # x1 = size of stack
+            # x2 = entry point
+            #
+            # Reserve space for the stack and align it
+            "eb 63 21 cb"  # sub x11, sp, x1
+            "6b ed 7c 92"  # and x11, x11, #~0xF
+            "7f 01 00 91"  # mov sp, x11
+            # Copy the stack over (dst=x3, src=x0, count=x1)
+            "e3 03 00 91"  # mov x3, sp
+            "a1 00 00 b4"  # loop: cbz x1, end
+            "06 14 40 38"  # ldrb w6, [x0], #1
+            "66 14 00 38"  # strb w6, [x3], #1
+            "21 04 00 f1"  # subs x1, x1, #1
+            "81 ff ff 54"  # b.ne loop
+            # Jump to the entry point
+            "40 00 1f d6"  # end: br x2
+        )
+    else:
+        raise NotImplementedError(f"Unsupported architecture: {machine}")
+
+    buf = libc.mmap(0, len(asm), PROT.WRITE, MAP.PRIVATE | MAP.ANONYMOUS, -1, 0)
+    ptr = ctypes.cast(buf, ctypes.POINTER(ctypes.c_char * len(asm)))
+    ptr.contents[:] = asm
+    libc.mprotect(buf, len(asm), PROT.READ | PROT.EXEC)
+
+    loader = ctypes.cast(
+        buf,
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_char_p,
+            ctypes.c_size_t,
+            ctypes.c_size_t,
+        ),
+    )
+
+    if log.level <= logging.DEBUG:
+        loader = _ctype_log(loader)
+
+
+def _ctype_log(function: Callable) -> Callable:
+    def fmt(arg: Any) -> str:
+        if isinstance(arg, enum.Enum):
+            return f"{arg.__class__.__name__}_{arg.name}"
+        if isinstance(arg, int):
+            return hex(arg)
+        return repr(arg)
+
+    def wrapper(*args, **kwargs) -> Any:
+        result = function(*args, **kwargs)
+        arg_str = ", ".join(f"{fmt(arg)}" for arg in args)
+        log.debug("%s(%s) -> %s", function.__name__, arg_str, fmt(result))
+        return result
+
+    return wrapper
+
+
+def load(
+    path: Path, argv: list[str] | None = None, envp: dict[str, str] | None = None, *, libc: str | None = None
+) -> None:
+    """Load an ELF executable into memory and execute it.
+
+    Args:
+        path: The path to the ELF executable to load.
+        argv: The arguments to pass to the executable.
+        envp: The environment variables to pass to the executable.
+        libc: The optional path to the libc library to use.
+    """
+    if os.name != "posix" and sys.platform != "darwin":
+        raise TypeError("This loader only supports POSIX systems")
+
+    _setup_libc(libc)
+    _setup_loader()
+
+    argv = argv or []
+    envp = envp or {}
+
+    with path.open("rb") as fh:
+        elf = ELF(fh)
+
+        if elf.dynamic:
+            raise NotImplementedError("Dynamic loading is not yet implemented, only static executables are supported")
+
+        segments = elf.segments.by_type(PT.LOAD)
+        base = truncate(min([s.virtual_address for s in segments]))
+
+        for segment in segments:
+            map_segment(segment, MAP.FIXED | MAP.PRIVATE | MAP.ANONYMOUS)
+
+        stack = create_stack(elf, base, argv, envp)
+        loader(stack, len(stack), elf.header.e_entry)
+
+
+def map_segment(segment: Segment, flags: MAP) -> None:
+    """Map segment to memory.
+
+    Args:
+        segment: The segment to map.
+        flags: The flags to pass to mmap.
+    """
+    alignment = segment.alignment - 1
+
+    offset = segment.virtual_address & alignment
+    start = truncate(segment.virtual_address, alignment)
+    size = round(segment.memory_size + offset, alignment)
+
+    pointer = libc.mmap(start, size, PROT.WRITE, flags, -1, 0)
+
+    if pointer == 0xFFFF_FFFF_FFFF_FFFF:
+        raise ValueError("mmap failed to allocate memory.")
+
+    ptr = ctypes.cast(pointer, ctypes.POINTER(ctypes.c_char * size))
+    ptr.contents[offset : offset + segment.size] = segment.data
+
+    pflags = PROT.from_PF(segment.flags)
+    libc.mprotect(pointer, size, pflags)
+
+
+def create_stack(elf: ELF, base: int, argv: list[str], envp: dict[str, str]) -> bytes:
+    """Create the initial process stack for the executable.
+
+    The stack is created according to the System V AMD64 ABI specification, which can be found here:
+    https://refspecs.linuxfoundation.org/elf/x86_64-abi-0.99.pdf
+
+    Args:
+        elf: The ELF executable to create the stack for.
+        base: The base address of the executable.
+        argv: The arguments to pass to the executable.
+        envp: The environment variables to pass to the executable.
+    """
+    stack = BytesIO()
+
+    stack.write(struct.pack("<Q", len(argv)))
+
+    c_argv = (ctypes.c_char_p * len(argv))(*[s.encode() for s in argv])
+    stack.write(c_argv)
+    stack.write(ctypes.c_void_p())
+
+    c_envp = (ctypes.c_char_p * len(envp))(*[f"{k}={v}".encode() for k, v in envp.items()])
+    stack.write(c_envp)
+    stack.write(ctypes.c_void_p())
+
+    copy_aux(stack, AT.SYSINFO_EHDR)
+    copy_aux(stack, AT.HWCAP)
+    copy_aux(stack, AT.PAGESZ)
+    copy_aux(stack, AT.CLKTCK)
+    write_aux(stack, AT.PHDR, ctypes.c_size_t(base + elf.header.e_phoff))
+    write_aux(stack, AT.PHENT, ctypes.c_size_t(elf.header.e_phentsize))
+    write_aux(stack, AT.PHNUM, ctypes.c_size_t(elf.header.e_phnum))
+    write_aux(stack, AT.BASE, ctypes.c_size_t(0))
+    write_aux(stack, AT.FLAGS, ctypes.c_size_t(0))
+    write_aux(stack, AT.ENTRY, ctypes.c_size_t(elf.header.e_entry))
+    copy_aux(stack, AT.UID)
+    copy_aux(stack, AT.EUID)
+    copy_aux(stack, AT.GID)
+    copy_aux(stack, AT.EGID)
+    copy_aux(stack, AT.SECURE)
+    write_aux(stack, AT.RANDOM, ctypes.c_char_p(random.randbytes(16)))
+    write_aux(stack, AT.EXECFN, ctypes.c_char_p(c_argv[0]))
+    write_aux(stack, AT.PLATFORM, ctypes.c_char_p(platform.machine().encode() + b"\x00"))
+    write_aux(stack, AT.NULL, ctypes.c_size_t(0))
+
+    return stack.getvalue()
+
+
+def copy_aux(fh: BinaryIO, type: AT) -> None:
+    """Copy auxiliary value from the host process to the stack."""
+    write_aux(fh, type, ctypes.c_uint64(libc.getauxval(type)))
+
+
+def write_aux(fh: BinaryIO, type: AT, value: bytes) -> None:
+    """Write auxiliary value to the stack."""
+    fh.write(struct.pack("<Q", type))
+    fh.write(value)
+
+
+def truncate(value: int, align: int = ALIGN) -> int:
+    """Truncate value to alignment."""
+    return value & ~(align)
+
+
+def round(value: int, align: int = ALIGN) -> int:
+    """Round value up to alignment."""
+    return truncate(value + align, align)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("path", type=Path, help="the ELF executable to load")
+    parser.add_argument("--libc", type=str, help="the path of the libc library to use")
+    parser.add_argument(
+        "--env", "-e", type=json.loads, default=None, help="environment variables to pass to the executable"
+    )
+    parser.add_argument("-v", "--verbose", action="store_true", help="enable verbose logging")
+    args, rest = parser.parse_known_args()
+
+    if args.verbose:
+        logging.basicConfig(format="%(asctime)s %(levelname)s %(message)s")
+        log.setLevel(logging.DEBUG)
+
+    arguments = [args.path.name, *rest]
+    load(args.path, arguments, args.env, libc=args.libc)
+
+
+if __name__ == "__main__":
+    main()

tests/_data/elf/hello_world.c

@@ -0,0 +1,7 @@
+// gcc hello_world.c -Wl,-Ttext-segment=0x1000000 [-static] -o hello_world.bin
+#include <stdio.h>
+
+int main() {
+    setbuf(stdout, NULL);
+    printf("Kusjes van SRT <3\n");
+}

tests/elf/test_dump.py

@@ -15,8 +15,8 @@
 @pytest.mark.parametrize(
     "path",
     [
-        "_data/elf/hello_world.out",
-        "_data/elf/hello_world.stripped.out",
+        "_data/elf/hello_world-x86_64.bin",
+        "_data/elf/hello_world-x86_64-stripped.bin",
     ],
 )
 def test_dump(path: str, tmp_path: Path) -> None: