#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Bytes-only dynamic/relocation contract layered on Phase 1.0AG."""

from __future__ import annotations

from dataclasses import dataclass
import struct

from phase10ag_bounded_elf import (
    ELF_HEADER, ElfContractError, PF_W, assess_elf,
)


SECTION_HEADER = struct.Struct("<IIQQQQIIQQ")
DYNAMIC_ENTRY = struct.Struct("<qQ")
RELA_ENTRY = struct.Struct("<QQq")
SHT_STRTAB = 3
SHT_RELA = 4
SHT_DYNAMIC = 6
DT_NULL = 0
DT_NEEDED = 1
R_X86_64_GLOB_DAT = 6
R_X86_64_RELATIVE = 8
MAX_DYNAMIC_ENTRIES = 256
MAX_RELOCATION_SECTIONS = 4
MAX_RELOCATIONS = 4096
MAX_NEEDED = 16
MAX_NEEDED_NAME = 128

ALLOWED_MODULES = frozenset({
    "libSceAudioOut.sprx", "libSceLibcInternal.sprx", "libScePad.sprx",
    "libSceSystemService.sprx", "libSceUserService.sprx",
    "libSceVideoOut.sprx", "libkernel_web.sprx",
})


class DynamicContractError(ValueError):
    """Dynamic metadata exceeds or violates the admitted subset."""


@dataclass(frozen=True)
class DynamicAssessment:
    needed: tuple[str, ...]
    relative_count: int
    glob_dat_count: int
    total_relocations: int
    relocation_sections: int
    loader_applied_types: tuple[int, ...] = (R_X86_64_RELATIVE,)
    crt_applied_types: tuple[int, ...] = (R_X86_64_GLOB_DAT,)
    unknown_relocation_types: tuple[int, ...] = ()
    target_mapping_performed: bool = False
    runtime_module_loading_performed: bool = False


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


def _sections(payload: bytes) -> list[tuple[int, ...]]:
    header = ELF_HEADER.unpack_from(payload)
    shoff, shentsize, shnum = header[6], header[11], header[12]
    if shnum == 0 or shentsize != SECTION_HEADER.size:
        raise DynamicContractError("section table is required")
    _range(shoff, shnum * shentsize, len(payload), "section table")
    return [SECTION_HEADER.unpack_from(payload, shoff + index * shentsize)
            for index in range(shnum)]


def _needed(payload: bytes, sections: list[tuple[int, ...]]) -> tuple[str, ...]:
    dynamic_sections = [section for section in sections if section[1] == SHT_DYNAMIC]
    if len(dynamic_sections) != 1:
        raise DynamicContractError("exactly one dynamic section is required")
    section = dynamic_sections[0]
    offset, size, link, entsize = section[4], section[5], section[6], section[9]
    if entsize != DYNAMIC_ENTRY.size or size % entsize or size // entsize > MAX_DYNAMIC_ENTRIES:
        raise DynamicContractError("dynamic table sizing is invalid")
    _range(offset, size, len(payload), "dynamic table")
    if link >= len(sections) or sections[link][1] != SHT_STRTAB:
        raise DynamicContractError("dynamic string table link is invalid")
    strings = sections[link]
    string_offset, string_size = strings[4], strings[5]
    _range(string_offset, string_size, len(payload), "dynamic string table")
    string_data = payload[string_offset:string_offset + string_size]
    names: list[str] = []
    terminated = False
    for position in range(offset, offset + size, entsize):
        tag, value = DYNAMIC_ENTRY.unpack_from(payload, position)
        if terminated:
            if tag != DT_NULL or value != 0:
                raise DynamicContractError("nonzero dynamic data follows DT_NULL")
            continue
        if tag == DT_NULL:
            terminated = True
            continue
        if tag != DT_NEEDED:
            continue
        if len(names) >= MAX_NEEDED or value >= len(string_data):
            raise DynamicContractError("DT_NEEDED count or offset is invalid")
        end = string_data.find(b"\0", value, min(len(string_data), value + MAX_NEEDED_NAME + 1))
        if end < 0:
            raise DynamicContractError("DT_NEEDED name is unterminated or too long")
        try:
            name = string_data[value:end].decode("ascii")
        except UnicodeDecodeError as error:
            raise DynamicContractError("DT_NEEDED name is not ASCII") from error
        if name not in ALLOWED_MODULES or name in names:
            raise DynamicContractError("DT_NEEDED module is unknown or duplicated")
        names.append(name)
    if not terminated:
        raise DynamicContractError("dynamic table has no DT_NULL terminator")
    return tuple(names)


def assess_dynamic(payload: bytes, expected_sha256: str,
                   expected_needed: tuple[str, ...]) -> DynamicAssessment:
    """Validate exact dynamic dependencies and the bounded relocation split."""
    try:
        elf = assess_elf(payload, expected_sha256)
    except ElfContractError as error:
        raise DynamicContractError("base ELF admission failed") from error
    if not isinstance(expected_needed, tuple) or not 1 <= len(expected_needed) <= MAX_NEEDED \
            or any(not isinstance(item, str) for item in expected_needed) \
            or len(set(expected_needed)) != len(expected_needed) \
            or any(item not in ALLOWED_MODULES for item in expected_needed):
        raise DynamicContractError("expected module inventory is invalid")
    sections = _sections(payload)
    needed = _needed(payload, sections)
    if needed != expected_needed:
        raise DynamicContractError("DT_NEEDED inventory or order differs")

    writable = [(item.virtual_address, item.virtual_address + item.memory_size)
                for item in elf.load_segments if item.flags & PF_W]
    image = [(item.virtual_address, item.virtual_address + item.memory_size)
             for item in elf.load_segments]
    relative = 0
    glob_dat = 0
    total = 0
    relocation_sections = 0
    for section in sections:
        if section[1] != SHT_RELA:
            continue
        relocation_sections += 1
        if relocation_sections > MAX_RELOCATION_SECTIONS:
            raise DynamicContractError("too many relocation sections")
        offset, size, entsize = section[4], section[5], section[9]
        if entsize != RELA_ENTRY.size or size % entsize:
            raise DynamicContractError("relocation table sizing is invalid")
        _range(offset, size, len(payload), "relocation table")
        for position in range(offset, offset + size, entsize):
            total += 1
            if total > MAX_RELOCATIONS:
                raise DynamicContractError("relocation count is exceeded")
            target, info, addend = RELA_ENTRY.unpack_from(payload, position)
            relocation_type = info & 0xffffffff
            symbol = info >> 32
            if target % 8 or not any(start <= target and target + 8 <= end
                                     for start, end in writable):
                raise DynamicContractError("relocation target is not aligned RW memory")
            if relocation_type == R_X86_64_RELATIVE:
                if symbol != 0 or addend < 0 or not any(start <= addend < end
                                                       for start, end in image):
                    raise DynamicContractError("relative relocation is invalid")
                relative += 1
            elif relocation_type == R_X86_64_GLOB_DAT:
                if symbol == 0 or addend != 0:
                    raise DynamicContractError("GLOB_DAT relocation is invalid")
                glob_dat += 1
            else:
                raise DynamicContractError("relocation type is outside the contract")
    if relocation_sections == 0 or relative == 0:
        raise DynamicContractError("relative relocation closure is absent")
    return DynamicAssessment(needed, relative, glob_dat, total,
                             relocation_sections)