remote play pup decrypter

#!/usr/bin/env python3

import struct
import sys
from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto.PublicKey import ECC
from Crypto.Signature import DSS

update_file = sys.argv[1]
VALIDATE_SIG = True

PSP_KEY = bytes.fromhex("3515a88f33557df133fbf208d63b0aaf")

OTA_PKG_PUB_PEM = """
-----BEGIN PUBLIC KEY-----
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAElbdJ1xORLXvkIske+h4rUAM/3TKq
KO3PpS/CmHKqICHlGyd2LGnPf/A2+2E+zKdfw1shWsoN5Le/oIbyL0coGQ==
-----END PUBLIC KEY-----
""".strip()

ota_pkg_pub = ECC.import_key(OTA_PKG_PUB_PEM)

# Get PUP file
with open(update_file, "rb") as fin:
    pupdata = fin.read()

# Validate file checksum
pupdata, file_checksum = pupdata[:-32], pupdata[-32:]
h = SHA256.new(pupdata)
assert h.digest() == file_checksum

# Parse plaintext header
header_to_verify = pupdata[:64]
magic, pupdata = pupdata[:4], pupdata[4:]
assert magic == b"DWCP"

header, pupdata = pupdata[: 4 + 8 + 8 + 4 + 4], pupdata[4 + 8 + 8 + 4 + 4 :]
header_version, file_size, body_size, attribute, reserved = struct.unpack(
    ">IQQII", header
)
print(
    f"header_version={header_version}, file_size={file_size:#x}, body_size={body_size:#x}, attribute={attribute}, reserved={reserved}"
)
header2, pupdata = pupdata[:16], pupdata[16:]
v_major, v_minor, v_patch, v_null = struct.unpack(">IIII", header2)
assert v_null == 0
print(f"version: v{v_major}.{v_minor}.{v_patch}")
iv, pupdata = pupdata[:16], pupdata[16:]

# Validate sizes
header_size = file_size - body_size
assert header_size >= 208
assert (header_size - 128 - 32) % 48 == 0

# Decrypt metadata block
aes = AES.new(key=PSP_KEY, mode=AES.MODE_CBC, iv=iv)
encrypted_metadata, pupdata = (
    pupdata[: header_size - 64 - 32 - 64],
    pupdata[header_size - 64 - 32 - 64 :],
)
metadata_block = aes.decrypt(encrypted_metadata)

# Parse metadata block
assert len(metadata_block) % 48 == 0
metadata = []
while len(metadata_block) > 0:
    key, iv, tag, metadata_block = (
        metadata_block[0:16],
        metadata_block[16:32],
        metadata_block[32:48],
        metadata_block[48:],
    )
    metadata.append((key, iv, tag))
print(f"File pages: {len(metadata)}")

# Validate header signature
verifier = DSS.new(ota_pkg_pub, mode="fips-186-3")
h = SHA256.new(header_to_verify + encrypted_metadata)
header_signature, pupdata = pupdata[:64], pupdata[64:]
try:
    verifier.verify(h, header_signature)
except ValueError as e:
    if VALIDATE_SIG:
        raise e
    else:
        print("Warning: Header signature invalid")

# Decrypt and verify data
enc_data = pupdata
decdata = b""
for idx, (key, iv, tag) in enumerate(metadata):
    page = enc_data[idx * 0x100000 : (idx + 1) * 0x100000]
    aes = AES.new(key=key, nonce=iv, mode=AES.MODE_GCM)

    m = aes.decrypt_and_verify(page, tag)
    decdata += m

# Save decrypted file
with open(f"{update_file}.jar", "wb") as fout:
    fout.write(decdata)