dvpl-engine 0.1.3

#!/usr/bin/env python3
"""DVPL file format converter for World of Tanks Blitz"""

from __future__ import annotations

from argparse import ArgumentParser
from pathlib import Path
from struct import pack
from struct import unpack
from sys import exit as sys_exit
from sys import stderr
from typing import Final
from zlib import crc32

from lz4.block import compress
from lz4.block import decompress


MAGIC: Final = b"DVPL"
FOOTER_SIZE: Final = 20
FOOTER_FMT: Final = "<IIII4s"  # original_size, compressed_size, crc32, compression_type, magic

COMP_NONE: Final = 0
COMP_LZ4: Final = 1
COMP_LZ4_HC: Final = 2


class DvplError(ValueError):
    """Base error for all DVPL operations"""


class TooSmallError(DvplError):
    """Input shorter than 20 bytes (footer size)"""


class BadMagicError(DvplError):
    """Footer magic does not match b"DVPL" """


class SizeMismatchError(DvplError):
    """Payload length disagrees with footer"""


class CrcMismatchError(DvplError):
    """CRC32 of payload does not match footer checksum"""


class DecompressedSizeMismatchError(DvplError):
    """Decompressed output length disagrees with footer"""


class UnknownCompressionError(DvplError):
    """Unrecognized compression type value"""


def decode(data: bytes) -> bytes:
    """Decode a DVPL-wrapped blob, verifying integrity"""
    if len(data) < FOOTER_SIZE:
        raise TooSmallError(f"File too small ({len(data)} bytes)")

    footer = data[-FOOTER_SIZE:]
    original_size, compressed_size, checksum, comp_type, magic = unpack(FOOTER_FMT, footer)

    if magic != MAGIC:
        raise BadMagicError(f"Bad magic: expected DVPL, got {magic!r}")

    payload = data[:-FOOTER_SIZE]
    if len(payload) != compressed_size:
        raise SizeMismatchError(f"Size mismatch: footer says {compressed_size}, got {len(payload)}")

    calculated_crc = crc32(payload) & 0xFFFFFFFF
    if calculated_crc != checksum:
        raise CrcMismatchError(
            f"CRC32 mismatch: expected {checksum:#010x}, got {calculated_crc:#010x}"
        )

    if comp_type == COMP_NONE:
        result = payload
    elif comp_type in (COMP_LZ4, COMP_LZ4_HC):
        result = decompress(payload, uncompressed_size=original_size)
    else:
        raise UnknownCompressionError(f"Unknown compression type: {comp_type}")

    if len(result) != original_size:
        raise DecompressedSizeMismatchError(
            f"Decompressed size mismatch: expected {original_size}, got {len(result)}"
        )

    return result


def encode(data: bytes, comp_type: int = COMP_LZ4_HC) -> bytes:
    """Encode raw data into DVPL format"""
    if comp_type == COMP_NONE:
        payload = data
    elif comp_type == COMP_LZ4:
        payload = compress(data, store_size=False)
    elif comp_type == COMP_LZ4_HC:
        payload = compress(data, mode="high_compression", store_size=False)
    else:
        raise UnknownCompressionError(f"Unknown compression type: {comp_type}")

    checksum = crc32(payload) & 0xFFFFFFFF
    footer = pack(FOOTER_FMT, len(data), len(payload), checksum, comp_type, MAGIC)

    return payload + footer


def _convert_file(path: Path, *, to_dvpl: bool, comp_type: int, output_dir: Path | None) -> Path:
    """Convert a single file, return output path"""
    data = path.read_bytes()

    if to_dvpl:
        result = encode(data, comp_type)
        out_path = (output_dir or path.parent) / (path.name + ".dvpl")
    else:
        result = decode(data)
        stem = (
            path.name.removesuffix(".dvpl")
            if path.name.endswith(".dvpl")
            else path.name + ".decoded"
        )
        out_path = (output_dir or path.parent) / stem

    out_path.write_bytes(result)
    return out_path


def main() -> None:
    """CLI entry point for DVPL conversion"""
    parser = ArgumentParser(description="Convert DVPL files (World of Tanks Blitz)")
    parser.add_argument("files", nargs="+", type=Path, help="Files to convert")
    parser.add_argument(
        "-e",
        "--encode",
        action="store_true",
        help="Encode to DVPL (default is decode)",
    )
    parser.add_argument(
        "-c",
        "--compression",
        type=int,
        default=COMP_LZ4_HC,
        choices=[0, 1, 2],
        help="Compression type for encoding: 0=none, 1=lz4, 2=lz4-hc (default: 2)",
    )
    parser.add_argument("-o", "--output-dir", type=Path, help="Output directory")
    args = parser.parse_args()

    if args.output_dir:
        args.output_dir.mkdir(parents=True, exist_ok=True)

    errors = 0
    for path in args.files:
        if not path.is_file():
            print(f"Skip: {path} (not a file)", file=stderr)
            errors += 1
            continue

        try:
            out = _convert_file(
                path,
                to_dvpl=args.encode,
                comp_type=args.compression,
                output_dir=args.output_dir,
            )
            print(f"{'Encoded' if args.encode else 'Decoded'}: {path} -> {out}")

        except Exception as e:
            print(f"Error: {path}: {e}", file=stderr)
            errors += 1

    if errors:
        sys_exit(1)


if __name__ == "__main__":
    main()