# SPDX-FileCopyrightText: 2026 Tomasz Kramkowski <tomasz@kramkow.ski>
# SPDX-License-Identifier: GPL-3.0-or-later

"""Compile C/assembly hooks into one Linux x86-64 ELF patch payload."""

import struct
import subprocess
from pathlib import Path
from typing import BinaryIO

from capstone import CS_ARCH_X86, CS_MODE_64, Cs
from elftools.construct import Container
from elftools.elf.constants import P_FLAGS
from elftools.elf.elffile import ELFFile
from elftools.elf.relocation import RelocationSection
from elftools.elf.sections import SymbolTableSection
from pydantic import BaseModel, ConfigDict, Field


class Hook(BaseModel):
    model_config = ConfigDict(strict=True, frozen=True, extra="forbid")

    address: int = Field(ge=0)
    length: int = Field(ge=5)
    target: str = Field(min_length=1)


class CompiledPayload:
    """Position-independent code with resolved entry addresses inside it."""

    code: bytes
    entries: dict[str, int]

    def __init__(self, source: BinaryIO, targets: list[str]):
        elf = ELFFile(source)
        for section in elf.iter_sections():
            if (
                section["sh_flags"] & 2
                and section["sh_size"]
                and section.name != ".payload"
            ):
                raise ValueError(f"unsupported allocated section: {section.name}")
            # PC32 (2) PLT32 (4) (PC-relative)
            if isinstance(section, RelocationSection):
                if any(
                    r["r_info_type"] not in (2, 4) for r in section.iter_relocations()
                ):
                    raise ValueError("helper contains a non-relative relocation")
        payload = elf.get_section_by_name(".payload")
        if payload is None:
            raise ValueError("missing compiled payload")
        self.code = payload.data()
        symbols = elf.get_section_by_name(".symtab")
        if not isinstance(symbols, SymbolTableSection):
            raise ValueError("missing compiled symbol table")
        payload_index = elf.get_section_index(".payload")
        address = payload["sh_addr"]
        self.entries = {}
        for target in targets:
            matches = symbols.get_symbol_by_name(target) or []
            if len(matches) != 1 or matches[0]["st_shndx"] != payload_index:
                raise ValueError(
                    f"hook target is not a unique payload symbol: {target}"
                )
            entry = matches[0]["st_value"]
            if not address <= entry < address + len(self.code):
                raise ValueError(
                    f"hook target is outside the compiled payload: {target}"
                )
            self.entries[target] = entry


def align(value: int, boundary: int) -> int:
    return (value + boundary - 1) & -boundary


def run(*args: str) -> str:
    return subprocess.check_output(args, text=True)


def resolve_hook_sites(elf: ELFFile, hooks: list[Hook]) -> list[tuple[Hook, int]]:
    """Find file offsets and verify that each hook covers whole instructions."""
    source = elf.stream
    sites = []
    disassembler = Cs(CS_ARCH_X86, CS_MODE_64)
    for hook in hooks:
        offset = next(elf.address_offsets(hook.address, hook.length))
        source.seek(offset)
        instructions = disassembler.disasm(source.read(hook.length), hook.address)
        if sum(i.size for i in instructions) != hook.length:
            raise ValueError("hook length cuts an instruction or contains invalid code")
        sites.append((hook, offset))
    return sites


def compile_code(
    sources: list[Path],
    imports: dict[str, int],
    targets: list[str],
    code_va: int,
    output: Path,
) -> CompiledPayload:
    """Compile and link the patch sources at their destination address."""
    output.mkdir(parents=True, exist_ok=False)
    objects = []
    flags = [
        "-Os",
        "-ffreestanding",
        "-fPIE",
        "-fvisibility=hidden",
        "-fno-stack-protector",
        "-fno-strict-aliasing",
        "-mno-red-zone",
        "-fno-asynchronous-unwind-tables",
        "-fno-unwind-tables",
        "-fcf-protection=none",
        "-Wall",
        "-Wextra",
    ]
    for number, source_path in enumerate(sources):
        obj = output / f"{number}.o"
        source_flags = ["-std=c17"] if source_path.suffix == ".c" else []
        run("cc", *flags, *source_flags, "-c", str(source_path), "-o", str(obj))
        objects.append(str(obj))

    linker_script = output / "layout.ld"
    linker_script.write_text(
        "PHDRS { code PT_LOAD FLAGS(5); }\n"  # PF_R | PF_X: readable/executable.
        f"SECTIONS {{ . = {code_va:#x};\n"
        " .payload : { *(.text.entry) *(.text*) *(.rodata*) } :code\n"
        " /DISCARD/ : { *(.comment) *(.note.GNU-stack) }\n}\n"
    )
    import_args = [f"--defsym={name}={value:#x}" for name, value in imports.items()]
    linked = output / "helper.elf"
    run(
        "ld",
        "--build-id=none",
        "--emit-relocs",
        f"--script={linker_script}",
        *import_args,
        f"--output={linked}",
        *objects,
    )
    with linked.open("rb") as compiled:
        return CompiledPayload(compiled, targets)


def make_hook_edits(
    sites: list[tuple[Hook, int]], entries: dict[str, int]
) -> list[tuple[int, bytes]]:
    """Redirect each hook site to its compiled entry with a relative jump."""
    edits = []
    for hook, offset in sites:
        displacement = entries[hook.target] - hook.address - 5
        if not -(1 << 31) <= displacement < 1 << 31:
            raise ValueError(
                f"hook target is outside relative jump range: {hook.target}"
            )
        jump = b"\xe9" + struct.pack("<i", displacement) + b"\x90" * (hook.length - 5)
        edits.append((offset, jump))
    return edits


def compile_segment(
    elf: ELFFile,
    sources: list[Path],
    imports: dict[str, int],
    targets: list[str],
    output: Path,
) -> tuple[CompiledPayload, list[tuple[int, bytes]]]:
    """Place the compiled code and enlarged program-header table in a new segment."""
    phdr = elf.structs.Elf_Phdr
    headers = [segment.header for segment in elf.iter_segments()]
    # Place the new segment after all existing loadable mappings.
    load_end = max(
        header["p_vaddr"] + header["p_memsz"]
        for header in headers
        if header["p_type"] == "PT_LOAD"
    )
    segment_va = align(load_end, 4096)
    segment_offset = align(elf.stream_len, 4096)
    table_size = (len(headers) + 1) * phdr.sizeof()
    code_delta = align(table_size, 16)
    code_va = segment_va + code_delta
    compiled = compile_code(sources, imports, targets, code_va, output)

    # Relocate the enlarged program-header table into the new segment.
    table_header = next(header for header in headers if header["p_type"] == "PT_PHDR")
    table_header.update(
        p_flags=P_FLAGS.PF_R,
        p_offset=segment_offset,
        p_vaddr=segment_va,
        p_paddr=segment_va,
        p_filesz=table_size,
        p_memsz=table_size,
        p_align=8,
    )
    segment_size = code_delta + len(compiled.code)
    headers.append(
        Container(
            p_type="PT_LOAD",
            p_flags=P_FLAGS.PF_R | P_FLAGS.PF_X,
            p_offset=segment_offset,
            p_vaddr=segment_va,
            p_paddr=segment_va,
            p_filesz=segment_size,
            p_memsz=segment_size,
            p_align=4096,
        )
    )
    edits = [
        (0x20, struct.pack("<Q", segment_offset)),  # Elf64_Ehdr.e_phoff
        (0x38, struct.pack("<H", len(headers))),  # Elf64_Ehdr.e_phnum
        (segment_offset, b"".join(phdr.build(header) for header in headers)),
        (segment_offset + code_delta, compiled.code),
    ]
    return compiled, edits


def compile_hooks(
    source: BinaryIO,
    sources: list[Path],
    imports: dict[str, int],
    hooks: list[Hook],
    output: Path,
) -> list[tuple[int, bytes]]:
    elf = ELFFile(source)
    sites = resolve_hook_sites(elf, hooks)
    compiled, segment_edits = compile_segment(
        elf, sources, imports, [hook.target for hook in hooks], output
    )
    return make_hook_edits(sites, compiled.entries) + segment_edits