#!/usr/bin/env python3
"""Bounded Git pack inflation and delta resolution for quarantine scanning.

Pushed objects arrive zlib-compressed inside a packfile, so any content
screening must inflate the real payloads first.  This module unpacks one
self-contained pack under explicit object, size, and checksum bounds and
resolves every delta against bases inside the same pack.  Thin packs are
rejected outright: a delta against a base the broker cannot see would leave
part of the pushed content unscanned.
"""

from __future__ import annotations

import hashlib
import tempfile
import zlib
from typing import BinaryIO

from gitea_api_policy import PolicyError

OBJECT_TYPE_NAMES = {1: b"commit", 2: b"tree", 3: b"blob", 4: b"tag"}
OFS_DELTA = 6
REF_DELTA = 7
MAX_PACK_OBJECTS = 50_000
MAX_OBJECT_BYTES = 32 * 1024 * 1024
MAX_TOTAL_BYTES = 256 * 1024 * 1024
MAX_RESOLVE_PASSES = 64
SPOOL_MEMORY_LIMIT = 256 * 1024
READ_CHUNK = 64 * 1024


def _read_exact(stream: BinaryIO, size: int) -> bytes:
    value = stream.read(size)
    if len(value) != size:
        raise PolicyError("Git pack data ended early")
    return value


def _object_header(stream: BinaryIO) -> tuple[int, int]:
    """Parse one object's type and declared inflated size."""
    byte = _read_exact(stream, 1)[0]
    type_code = (byte >> 4) & 0x7
    size = byte & 0x0F
    shift = 4
    while byte & 0x80:
        if shift > 60:
            raise PolicyError("Git object header is invalid")
        byte = _read_exact(stream, 1)[0]
        size |= (byte & 0x7F) << shift
        shift += 7
    return type_code, size


def _delta_distance(stream: BinaryIO) -> int:
    """Parse one ofs-delta base distance."""
    byte = _read_exact(stream, 1)[0]
    value = byte & 0x7F
    while byte & 0x80:
        byte = _read_exact(stream, 1)[0]
        value = ((value + 1) << 7) | (byte & 0x7F)
        if value > 1 << 48:
            raise PolicyError("Git delta base distance is invalid")
    return value


def _inflate(stream: BinaryIO, declared: int) -> BinaryIO:
    """Inflate one object into a quarantine spool, bounded by its header."""
    spool = tempfile.SpooledTemporaryFile(  # noqa: SIM115 - caller owns returned spool
        max_size=SPOOL_MEMORY_LIMIT, dir="/tmp"
    )
    decompressor = zlib.decompressobj()
    produced = 0
    try:
        while not decompressor.eof:
            data = stream.read(READ_CHUNK)
            if not data:
                raise PolicyError("Git pack object data ended early")
            while True:
                value = decompressor.decompress(data, READ_CHUNK)
                produced += len(value)
                if produced > declared:
                    raise PolicyError("Git object size does not match its header")
                spool.write(value)
                data = decompressor.unconsumed_tail
                if decompressor.eof or not data:
                    break
        if decompressor.unused_data:
            stream.seek(-len(decompressor.unused_data), 1)
        if produced != declared:
            raise PolicyError("Git object size does not match its header")
        spool.seek(0)
        return spool
    except zlib.error as exc:
        spool.close()
        raise PolicyError("Git pack object data is corrupt") from exc
    except Exception:
        spool.close()
        raise


def _delta_size(delta: bytes, position: int) -> tuple[int, int]:
    size = 0
    shift = 0
    while True:
        if position >= len(delta) or shift > 60:
            raise PolicyError("Git delta header is invalid")
        byte = delta[position]
        position += 1
        size |= (byte & 0x7F) << shift
        shift += 7
        if not byte & 0x80:
            return size, position


def _apply_delta(base: bytes, delta: bytes) -> bytes:
    """Apply one Git delta and require exact declared sizes throughout."""
    source_size, position = _delta_size(delta, 0)
    target_size, position = _delta_size(delta, position)
    if source_size != len(base) or target_size > MAX_OBJECT_BYTES:
        raise PolicyError("Git delta sizes are invalid")
    parts: list[bytes] = []
    produced = 0
    while position < len(delta):
        opcode = delta[position]
        position += 1
        if opcode & 0x80:
            operands = [bit for bit in (1, 2, 4, 8, 16, 32, 64) if opcode & bit]
            if position + len(operands) > len(delta):
                raise PolicyError("Git delta copy is truncated")
            offset = 0
            size = 0
            for shift, bit in enumerate((1, 2, 4, 8)):
                if opcode & bit:
                    offset |= delta[position] << (8 * shift)
                    position += 1
            for shift, bit in enumerate((16, 32, 64)):
                if opcode & bit:
                    size |= delta[position] << (8 * shift)
                    position += 1
            size = size or 0x10000
            if offset + size > len(base):
                raise PolicyError("Git delta copy is out of bounds")
            parts.append(base[offset : offset + size])
            produced += size
        elif opcode:
            if position + opcode > len(delta):
                raise PolicyError("Git delta insert is truncated")
            parts.append(delta[position : position + opcode])
            position += opcode
            produced += opcode
        else:
            raise PolicyError("Git delta instruction is invalid")
        if produced > target_size:
            raise PolicyError("Git delta result exceeds its declared size")
    if produced != target_size:
        raise PolicyError("Git delta result size is invalid")
    return b"".join(parts)


def _payload_sha(type_code: int, payload: BinaryIO, size: int) -> bytes:
    digest = hashlib.sha1(OBJECT_TYPE_NAMES[type_code] + b" %d\x00" % size)
    payload.seek(0)
    while True:
        chunk = payload.read(READ_CHUNK)
        if not chunk:
            break
        digest.update(chunk)
    payload.seek(0)
    return digest.digest()


def _verify_checksum(stream: BinaryIO, start: int, end: int) -> None:
    trailer = _read_exact(stream, 20)
    stream.seek(start)
    digest = hashlib.sha1()
    remaining = end - start
    while remaining:
        chunk = stream.read(min(READ_CHUNK, remaining))
        digest.update(chunk)
        remaining -= len(chunk)
    if digest.digest() != trailer:
        raise PolicyError("Git pack checksum is invalid")
    stream.seek(end + 20)


class _Entry:
    """One pack entry: either resolved content or an unresolved delta."""

    def __init__(self, type_code: int, payload: BinaryIO, size: int, base):
        self.type_code = type_code
        self.payload = payload
        self.size = size
        self.base = base


def _resolve(entries: dict[int, _Entry], total: list[int]) -> None:
    """Resolve every delta against in-pack bases only, to a fixed point."""
    by_sha: dict[bytes, int] = {}
    for offset, entry in entries.items():
        if entry.base is None:
            by_sha.setdefault(
                _payload_sha(entry.type_code, entry.payload, entry.size), offset
            )
    for _ in range(MAX_RESOLVE_PASSES):
        pending = [item for item in entries.items() if item[1].base is not None]
        if not pending:
            return
        progressed = False
        for offset, entry in pending:
            base_key = entry.base
            base_offset = base_key if isinstance(base_key, int) else by_sha.get(base_key)
            base = entries.get(base_offset) if base_offset is not None else None
            if base is None or base.base is not None:
                if isinstance(base_key, int) and base is None:
                    raise PolicyError("Git delta base offset is invalid")
                continue
            base.payload.seek(0)
            entry.payload.seek(0)
            value = _apply_delta(base.payload.read(), entry.payload.read())
            total[0] += len(value)
            if total[0] > MAX_TOTAL_BYTES:
                raise PolicyError("Git pack contents exceed the safe size limit")
            entry.payload.close()
            spool = tempfile.SpooledTemporaryFile(  # noqa: SIM115 - entry owns spool
                max_size=SPOOL_MEMORY_LIMIT, dir="/tmp"
            )
            spool.write(value)
            spool.seek(0)
            entry.type_code = base.type_code
            entry.payload = spool
            entry.size = len(value)
            entry.base = None
            by_sha.setdefault(_payload_sha(entry.type_code, spool, entry.size), offset)
            progressed = True
        if not progressed:
            raise PolicyError("Git delta base is outside the pack; push full packs")
    raise PolicyError("Git delta chains exceed the safe resolution limit")


def unpack_objects(stream: BinaryIO) -> list[tuple[int, BinaryIO]]:
    """Inflate one self-contained pack and return every resolved payload."""
    start = stream.tell()
    header = _read_exact(stream, 12)
    version = int.from_bytes(header[4:8], "big")
    count = int.from_bytes(header[8:12], "big")
    if header[:4] != b"PACK" or version not in {2, 3}:
        raise PolicyError("Git pack header is invalid")
    if count > MAX_PACK_OBJECTS:
        raise PolicyError("Git pack contains too many objects")
    entries: dict[int, _Entry] = {}
    total = [0]
    try:
        for _ in range(count):
            offset = stream.tell() - start
            type_code, declared = _object_header(stream)
            if declared > MAX_OBJECT_BYTES:
                raise PolicyError("Git object exceeds the safe size limit")
            total[0] += declared
            if total[0] > MAX_TOTAL_BYTES:
                raise PolicyError("Git pack contents exceed the safe size limit")
            base: int | bytes | None = None
            if type_code == OFS_DELTA:
                distance = _delta_distance(stream)
                base = offset - distance
                if base < 0:
                    raise PolicyError("Git delta base offset is invalid")
            elif type_code == REF_DELTA:
                base = _read_exact(stream, 20)
            elif type_code not in OBJECT_TYPE_NAMES:
                raise PolicyError("Git pack object type is invalid")
            entries[offset] = _Entry(type_code, _inflate(stream, declared), declared, base)
        _verify_checksum(stream, start, stream.tell())
        _resolve(entries, total)
    except Exception:
        for entry in entries.values():
            entry.payload.close()
        raise
    for entry in entries.values():
        entry.payload.seek(0)
    return [
        (entries[offset].type_code, entries[offset].payload)
        for offset in sorted(entries)
    ]


def parse_loose_object(value: bytes) -> tuple[int, bytes]:
    """Inflate one quarantined loose object and return its typed payload."""
    if len(value) > MAX_OBJECT_BYTES:
        raise PolicyError("Git object exceeds the safe size limit")
    decompressor = zlib.decompressobj()
    try:
        inflated = decompressor.decompress(value, MAX_OBJECT_BYTES + 1)
    except zlib.error as exc:
        raise PolicyError("Git loose object data is corrupt") from exc
    if not decompressor.eof or decompressor.unused_data or len(inflated) > MAX_OBJECT_BYTES:
        raise PolicyError("Git loose object framing is invalid")
    kind, separator, remainder = inflated.partition(b"\x00")
    name, space, size = kind.partition(b" ")
    types = {label: code for code, label in OBJECT_TYPE_NAMES.items()}
    if (
        not separator
        or not space
        or name not in types
        or not size.isdigit()
        or int(size) != len(remainder)
    ):
        raise PolicyError("Git loose object header is invalid")
    return types[name], remainder