Mach-O code injection with LIEF and shell-factory

inject.py

#!/usr/bin/env python

# Script associated with the blog post: https://lief-project.github.io/blog/2022-05-08-macho/
# It demonstrates code injection with shell-factory and LIEF

import lief
import pathlib
from pathlib import Path

CWD = Path(__file__).parent

target_path = CWD / "_heapq.cpython-39-darwin.so"
shellcode = CWD / "lief_demo_darwin_arm64.bin"

SEGMENTS = dict()

def get_imagebase_offset(target: lief.MachO.Binary, shellcode: lief.MachO.Binary) -> int:
    """
    This function returns the offset of the IMAGEBASE variable within the target and
    after the injection
    """

    IMAGEBASE = shellcode.get_symbol("_IMAGEBASE")
    shellcode_segment = shellcode.segment_from_virtual_address(IMAGEBASE.value)

    new_segment = SEGMENTS.get(shellcode_segment)
    assert new_segment is not None
    # Offset relative of the segment in the **shellcode**
    rva = IMAGEBASE.value - shellcode_segment.virtual_address
    return new_segment.virtual_address + rva

def get_oep_offset(target: lief.MachO.Binary, shellcode: lief.MachO.Binary):
    """
    This function returns the new entrypoint of the shellcode within the target
    """
    ORIGINAL_EP = shellcode.get_symbol("_ORIGINAL_EP")
    shellcode_segment = shellcode.segment_from_virtual_address(ORIGINAL_EP.value)

    new_segment = SEGMENTS.get(shellcode_segment)
    assert new_segment is not None
    # Offset relative of the segment in the **shellcode**
    rva = ORIGINAL_EP.value - shellcode_segment.virtual_address
    return new_segment.virtual_address + rva


def insert_shellcode(target: lief.MachO.Binary, shellcode: lief.MachO.Binary):
    """
    This function copies the segments from the shellcode in the target
    """
    SKIP_LIST = {"__PAGEZERO", "__LINKEDIT"}
    def is_valid_segment(seg: lief.MachO.SegmentCommand):
        return seg.name not in SKIP_LIST

    IMAGEBASE   = shellcode.get_symbol("_IMAGEBASE")
    ORIGINAL_EP = shellcode.get_symbol("_ORIGINAL_EP")

    shellcode_ep_offset = shellcode.virtual_address_to_offset(shellcode.entrypoint)
    for segment in filter(is_valid_segment, shellcode.segments):
        print("Adding {:<10}: [0x{:08x}, 0x{:08x}] in {}".format(
              segment.name, segment.virtual_address, segment.virtual_address + segment.virtual_size,
              Path(target.name).name))
        new_seg_name = segment.name.replace("__", "")
        new_seg = lief.MachO.SegmentCommand(f"__L{new_seg_name}", list(segment.content))
        sadded = target.add(new_seg)
        print("-----> {:<10}: [0x{:08x}, 0x{:08x}]".format(sadded.name,
              sadded.virtual_address, sadded.virtual_address + sadded.virtual_size))

        # Copy the memory protections
        sadded.init_protection = segment.init_protection
        sadded.max_protection  = segment.max_protection
        SEGMENTS[segment] = sadded

    __LTEXT = target.get_segment("__LTEXT")

    address = shellcode_ep_offset + __LTEXT.virtual_address - target.imagebase
    address -= target.imagebase # Export addresses are relative to the imagebase
    print(f"RVA of the shellcode's entrypoint: 0x{address:x}")
    return address

def patch_export(target: lief.MachO.Binary, shellcode: lief.MachO.Binary, symbol: str, value: int):
    """
    This function redirects an export to another function. It works by patching the export trie.
    """
    dyldinfo: lief.MachO.DyldInfo = target.dyld_info

    exp: lief.MachO.ExportInfo
    for exp in dyldinfo.exports:
        name = exp.symbol.name
        if name != symbol:
            continue

        print("Patching '{}' 0x{:08x} --> 0x{:08x}".format(name, exp.address, value))
        original = exp.address
        exp.address = value
        return original

# Parse both the shellcode and the target
target: lief.MachO.Binary = lief.parse(target_path.as_posix())
shellcode = lief.parse(shellcode.as_posix())

# Insert the shellcode in the target and get the **new** entrypoint
# of the shellcode within the target
ep = insert_shellcode(target, shellcode)

imagebase_va = get_imagebase_offset(target, shellcode)
print(f"IMAGEBASE   in the patched binary: 0x{imagebase_va:08x}")
target.patch_address(imagebase_va, imagebase_va, 8)

# Get the shellcode's entrypoint within the target
oep_va = get_oep_offset(target, shellcode)
print(f"ORIGINAL_EP in the patched binary: 0x{oep_va:08x}")

original_value = patch_export(target, shellcode, "_PyInit__heapq", ep)
target.patch_address(oep_va, original_value, 8)

# Write back the modified library
output = CWD / "_heapq.cpython-39-darwin.so.patched"
target.write(output.as_posix())

# Check the layout of the written binary (mostly for code signing)
new = lief.parse(output.as_posix())
has_err, err = lief.MachO.check_layout(new)
if has_err:
    print(err)
else:
    print(f"{output} has a correct layout!")