#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Bytes-only bounded ELF64 admission contract for a future launcher.

No path, file, process, target, network or execution interface is present.
"""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import re
import struct


MAX_FILE_SIZE = 2 * 1024 * 1024
MAX_PROGRAM_HEADERS = 32
MAX_SECTION_HEADERS = 256
MAX_LOAD_SEGMENTS = 8
MAX_TOTAL_LOAD_MEMORY = 64 * 1024 * 1024
MAX_LOAD_SPAN = 128 * 1024 * 1024
MAX_ALIGNMENT = 2 * 1024 * 1024

ELF_HEADER = struct.Struct("<16sHHIQQQIHHHHHH")
PROGRAM_HEADER = struct.Struct("<IIQQQQQQ")
SECTION_HEADER_SIZE = 64
PT_LOAD = 1
PT_INTERP = 3
ET_DYN = 3
EM_X86_64 = 62
PF_X = 1
PF_W = 2
PF_R = 4
SHA256 = re.compile(r"^[0-9a-f]{64}$")


class ElfContractError(ValueError):
    """The supplied bytes are not an admitted bounded payload."""


@dataclass(frozen=True)
class LoadSegment:
    flags: int
    offset: int
    virtual_address: int
    file_size: int
    memory_size: int
    alignment: int


@dataclass(frozen=True)
class ElfAssessment:
    size: int
    sha256: str
    entry: int
    program_header_count: int
    load_segments: tuple[LoadSegment, ...]
    total_load_memory: int
    load_span: int
    pie: bool = True
    machine_x86_64: bool = True
    writable_executable_segment: bool = False
    interpreter_present: bool = False
    execution_performed: bool = False


def _bounded_range(offset: int, size: int, limit: int, label: str) -> None:
    if offset < 0 or size < 0 or offset > limit or size > limit - offset:
        raise ElfContractError(f"{label} range is outside supplied bytes")


def _power_of_two(value: int) -> bool:
    return value > 0 and value & (value - 1) == 0


def assess_elf(payload: bytes, expected_sha256: str) -> ElfAssessment:
    """Validate exact supplied bytes without opening or executing anything."""
    if not isinstance(payload, bytes) or not ELF_HEADER.size <= len(payload) <= MAX_FILE_SIZE:
        raise ElfContractError("payload size is outside the admission bound")
    if not isinstance(expected_sha256, str) or not SHA256.fullmatch(expected_sha256):
        raise ElfContractError("expected SHA-256 is invalid")
    digest = hashlib.sha256(payload).hexdigest()
    if digest != expected_sha256:
        raise ElfContractError("payload SHA-256 does not match")

    (ident, elf_type, machine, version, entry, phoff, shoff, flags,
     ehsize, phentsize, phnum, shentsize, shnum, shstrndx) = ELF_HEADER.unpack_from(payload)
    del flags, shstrndx
    if ident[:4] != b"\x7fELF" or ident[4] != 2 or ident[5] != 1 \
            or ident[6] != 1 or ident[7] not in {0, 9}:
        raise ElfContractError("ELF identity is unsupported")
    if elf_type != ET_DYN or machine != EM_X86_64 or version != 1:
        raise ElfContractError("ELF type, machine or version is unsupported")
    if ehsize != ELF_HEADER.size or phentsize != PROGRAM_HEADER.size \
            or not 1 <= phnum <= MAX_PROGRAM_HEADERS:
        raise ElfContractError("ELF header sizing is invalid")
    if phoff < ELF_HEADER.size:
        raise ElfContractError("program header table overlaps ELF header")
    _bounded_range(phoff, phnum * phentsize, len(payload), "program header table")
    if shnum == 0:
        if shoff != 0 or shentsize not in {0, SECTION_HEADER_SIZE}:
            raise ElfContractError("absent section table is inconsistent")
    else:
        if not 1 <= shnum <= MAX_SECTION_HEADERS or shentsize != SECTION_HEADER_SIZE:
            raise ElfContractError("section header sizing is invalid")
        _bounded_range(shoff, shnum * shentsize, len(payload), "section header table")

    loads: list[LoadSegment] = []
    interpreter = False
    for index in range(phnum):
        values = PROGRAM_HEADER.unpack_from(payload, phoff + index * phentsize)
        p_type, p_flags, p_offset, p_vaddr, _p_paddr, p_filesz, p_memsz, p_align = values
        if p_type == PT_INTERP:
            interpreter = True
        if p_type != PT_LOAD:
            continue
        if len(loads) >= MAX_LOAD_SEGMENTS or p_memsz == 0 or p_filesz > p_memsz:
            raise ElfContractError("load segment count or sizing is invalid")
        _bounded_range(p_offset, p_filesz, len(payload), "load segment file")
        if p_vaddr > (1 << 64) - 1 - p_memsz:
            raise ElfContractError("load segment address overflows")
        if p_flags & ~(PF_R | PF_W | PF_X) or p_flags & PF_W and p_flags & PF_X:
            raise ElfContractError("load segment permissions are invalid")
        if not _power_of_two(p_align) or p_align > MAX_ALIGNMENT \
                or p_offset % p_align != p_vaddr % p_align:
            raise ElfContractError("load segment alignment is invalid")
        loads.append(LoadSegment(p_flags, p_offset, p_vaddr, p_filesz,
                                 p_memsz, p_align))
    if interpreter:
        raise ElfContractError("interpreter segment is forbidden")
    if not loads:
        raise ElfContractError("ELF has no load segments")

    ordered = sorted(loads, key=lambda item: item.virtual_address)
    for previous, current in zip(ordered, ordered[1:]):
        if previous.virtual_address + previous.memory_size > current.virtual_address:
            raise ElfContractError("load segment virtual ranges overlap")
    total_memory = sum(item.memory_size for item in ordered)
    span = ordered[-1].virtual_address + ordered[-1].memory_size - ordered[0].virtual_address
    if total_memory > MAX_TOTAL_LOAD_MEMORY or span > MAX_LOAD_SPAN:
        raise ElfContractError("load memory budget is exceeded")
    executable = [item for item in ordered if item.flags & PF_X]
    if not executable or not any(item.virtual_address <= entry <
                                 item.virtual_address + item.memory_size
                                 for item in executable):
        raise ElfContractError("entry is not in an executable load segment")
    return ElfAssessment(len(payload), digest, entry, phnum, tuple(ordered),
                         total_memory, span)