Unpack YesSteveModel's proprietary format into zip format, for version < 1.2.0 only.

unpack_ysm.py

import hashlib
import zipfile
import zlib
from base64 import b64decode
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import IO, Literal, Optional

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


class JavaRandom:
    def __init__(self, seed: int):
        self.seed = seed

    @property
    def seed(self):
        return self._seed

    @seed.setter
    def seed(self, seed: int):
        self._seed = (seed ^ 0x5DEECE66D) & ((1 << 48) - 1)

    def next(self, bits: int):
        """
        Generate the next random number.
        As in Java, the general rule is that this method returns an int that
        is `bits` bits long, where each bit is nearly equally likely to be 0
        or 1.
        """

        if bits < 1:
            bits = 1
        elif bits > 32:
            bits = 32

        self._seed = (self._seed * 0x5DEECE66D + 0xB) & ((1 << 48) - 1)
        retval = self._seed >> (48 - bits)

        # Python and Java don't really agree on how ints work. This converts
        # the unsigned generated int into a signed int if necessary.
        if retval & (1 << 31):
            retval -= 1 << 32

        return retval

    def next_bytes(self, length: int):
        buf = bytearray(length)
        for i in range(0, length):
            if not i % 4:
                n = self.next_int()
            b = n & 0xFF
            buf[i] = b
            n >>= 8
        return bytes(buf)

    def next_int(self, bits: Optional[int] = None):
        """
        Return a random int in [0, `n`).
        If `n` is not supplied, a random 32-bit integer will be returned.
        """

        if bits is None:
            return self.next(32)

        if bits <= 0:
            raise ValueError("Argument must be positive!")

        # This tricky chunk of code comes straight from the Java spec. In
        # essence, the algorithm tends to have much better entropy in the
        # higher bits of the seed, so this little bundle of joy is used to try
        # to reject values which would be obviously biased. We do have an easy
        # out for power-of-two n, in which case we can call next directly.

        # Is this a power of two?
        if not (bits & (bits - 1)):
            return (bits * self.next(31)) >> 31

        bits = self.next(31)
        val = bits % bits
        while (bits - val + bits - 1) < 0:
            bits = self.next(31)
            val = bits % bits

        return val


@dataclass(frozen=True)
class YSModelFile:
    filename: str
    data: bytes

    @classmethod
    def from_encrypted(cls, reader: IO[bytes], version: Literal[1, 2]):
        return {
            1: cls._decrypt_buffer_v1,
            2: cls._decrypt_buffer_v2,
        }[version](reader)

    @classmethod
    def _decrypt(cls, buf: bytes, key: bytes, iv: bytes):
        """Implements AES/CBC/PKCS5Padding decryption."""
        decryptor = Cipher(
            algorithms.AES(key),
            modes.CBC(iv),
        ).decryptor()
        decrypted = decryptor.update(buf) + decryptor.finalize()
        return decrypted[: -decrypted[-1]]  # PKCS#5 padding

    @staticmethod
    def _inflate(buf: bytes):
        """Implements Java's Inflater.inflate."""
        return zlib.decompress(buf)

    @staticmethod
    def _kdf(buf: bytes):
        """Derive a key from the given buffer."""
        data_seed = int.from_bytes(hashlib.md5(buf).digest(), "big")
        random = JavaRandom(data_seed)
        return random.next_bytes(16)

    @classmethod
    def _decrypt_buffer_v1(cls, reader: IO[bytes]):
        file_name_length = int.from_bytes(reader.read(4), "big")
        file_name = reader.read(file_name_length).decode()

        file_length = int.from_bytes(reader.read(4), "big")

        key = reader.read(16)
        nonce = reader.read(16)
        file = reader.read(file_length)

        file = cls._decrypt(file, key, nonce)
        file = cls._inflate(file)
        return cls(file_name, file)

    @classmethod
    def _decrypt_buffer_v2(cls, reader: IO[bytes]):
        filename_length = int.from_bytes(reader.read(4), "big")
        filename = b64decode(reader.read(filename_length)).decode()

        file_read_length = int.from_bytes(reader.read(4), "big")

        encrypted_key_read_length = int.from_bytes(reader.read(4), "big")
        encrypted_key_read = reader.read(encrypted_key_read_length)

        nonce = reader.read(16)
        file_read = reader.read(file_read_length)
        key = cls._decrypt(encrypted_key_read, cls._kdf(file_read), nonce)

        file = cls._decrypt(file_read, key, nonce)
        file = cls._inflate(file)
        return cls(filename, file)


def read_file(path: Path):
    with path.open("rb") as fd:
        if fd.read(4) != b"YSGP":
            raise ValueError("invalid file header")

        version = int.from_bytes(fd.read(4), "big")
        if version not in (1, 2):
            raise ValueError(f"unsupported version {version}")

        checksum = fd.read(16)
        buf = fd.read()
        if checksum != (actual_sum := hashlib.md5(buf).digest()):
            raise ValueError(f"invalid checksum, {checksum=} {actual_sum=} disagree")

        reader = BytesIO(buf)
        while reader.tell() < len(buf):  # loop until EOF
            yield YSModelFile.from_encrypted(reader, version)
    return


if __name__ == "__main__":
    import sys

    program_name, *args = sys.argv

    if len(args) != 1:
        print(f"Usage: {program_name} <ysm files directory or file>")
        sys.exit(1)

    file_or_dir = Path(args[0])
    if file_or_dir.is_file():
        files = [file_or_dir]
    else:
        files = file_or_dir.glob("*.ysm")

    for file in files:
        zip_file = file.with_suffix(".zip")
        if zip_file.exists():
            zip_file.unlink()
        with zipfile.ZipFile(
            zip_file,
            "w",
            compression=zipfile.ZIP_DEFLATED,
            compresslevel=zlib.Z_BEST_COMPRESSION,
        ) as z:
            for model in read_file(file):
                z.writestr(model.filename, model.data)
        print(f"Converted {file} to {zip_file}")