Initial commit

LICENCE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2021 donmai-me
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

MANIFEST.in

@@ -0,0 +1,2 @@
+include README.md
+include LICENSE

README.md

@@ -0,0 +1,32 @@
+# WannaCRI
+A (WIP) Python library for parsing, extracting, and generating Criware's various audio and video file formats.
+If you're interested in reading more about USM, you can read my write-up about it [here](https://listed.to/@donmai/24921/criware-s-usm-format-part-1)
+
+Currently supports the following formats with more planned:
+* USM (encrypted and plaintext)
+    * Vp9
+    * h264 (in-progress)
+
+
+This library has the following requirements:
+
+A working FFmpeg and FFprobe installation. On Windows, you can download official ffmpeg and ffprobe binaries and place them on your path.
+
+This project also heavily uses the [ffmpeg-python](https://pypi.org/project/ffmpeg-python) wrapper.
+
+# Usage
+
+If installed, there should be a command-line tool available.
+
+For extracting USMs:
+
+`wannacri extractusm /path/to/usm/file/or/folder --key 0xKEYUSEDIFENCRYPTED`
+
+For creating USMs:
+
+`wannacri createusm /path/to/vp9/file --key 0xKEYIFYOUWANTTOENCRYPT`
+
+# Licence
+
+This is an open-sourced application licensed under the MIT License
+

pyproject.toml

@@ -0,0 +1,4 @@
+[build-system]
+requires = ["setuptools>=49", "wheel", "setuptools_scm[toml]>=6.0"]
+
+[tool.setuptools_scm]

requirements.txt

@@ -0,0 +1 @@
+ffmpeg-python~=0.2.0

setup.py

@@ -0,0 +1,34 @@
+from setuptools import setup
+
+with open("README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+setup(
+    name="WannaCRI",
+    description="Criware media formats library",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="donmai",
+    url="https://github.com/donmai-me/WannaCRI",
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.8",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Developers",
+        "Topic :: Games/Entertainment",
+    ],
+    packages=[
+        "wannacri",
+        "wannacri.usm",
+        "wannacri.usm.media",
+    ],
+    entry_points={
+        "console_scripts": ["wannacri=wannacri:main"],
+    },
+    python_requires="~=3.8",
+    use_scm_version=True,
+    setup_requires=["setuptools_scm"],
+    install_requires=["ffmpeg-python~=0.2.0"],
+)

wannacri/__init__.py

@@ -0,0 +1,8 @@
+from importlib.metadata import version, PackageNotFoundError
+from .wannacri import main
+
+try:
+    __version__ = version("wannacri")
+except PackageNotFoundError:
+    # Not installed
+    __version__ = "not installed"

wannacri/codec.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from enum import Enum, auto
+import ffmpeg
+
+
+class Sofdec2Codec(Enum):
+    PRIME = auto()  # MPEG2
+    H264 = auto()
+    VP9 = auto()
+
+    @staticmethod
+    def from_file(path: str, ffprobe_path: str = "ffprobe") -> Sofdec2Codec:
+        info = ffmpeg.probe(path, cmd=ffprobe_path)
+
+        if len(info.get("streams")) == 0:
+            raise ValueError("File has no videos streams.")
+
+        codec_name = info.get("streams")[0].get("codec_name")
+        if codec_name == "vp9":
+            if info.get("format").get("format_name") != "ivf":
+                raise ValueError("VP9 file must be stored as an ivf.")
+
+            return Sofdec2Codec.VP9
+        if codec_name == "h264":
+            # TODO: Check if we need to have extra checks on h264 bitstreams
+            return Sofdec2Codec.H264
+        if codec_name == "mpeg2video":
+            # TODO: Check if we need to have extra checks on h264 bitstreams
+            return Sofdec2Codec.PRIME
+
+        raise ValueError(f"Unknown codec {codec_name}")

wannacri/usm/__init__.py

@@ -0,0 +1,21 @@
+from .tools import (
+    chunk_size_and_padding,
+    generate_keys,
+    is_valid_chunk,
+    encrypt_video_packet,
+    decrypt_video_packet,
+    encrypt_audio_packet,
+    decrypt_audio_packet,
+    get_video_header_end_offset,
+    is_usm,
+)
+from .page import UsmPage, get_pages, pack_pages
+from .usm import Usm
+from .chunk import UsmChunk
+from .media import UsmMedia, UsmVideo, UsmAudio, GenericVideo, GenericAudio, Vp9
+from .types import OpMode, ArrayType, ElementType, PayloadType, ChunkType
+
+import logging
+from logging import NullHandler
+
+logging.getLogger(__name__).addHandler(NullHandler())

wannacri/usm/chunk.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+import logging
+from typing import List, Union, Callable
+
+from .types import ChunkType, PayloadType
+from .page import UsmPage, pack_pages, get_pages
+from .tools import bytes_to_hex, is_valid_chunk
+
+
+class UsmChunk:
+    def __init__(
+        self,
+        chunk_type: ChunkType,
+        payload_type: PayloadType,
+        payload: Union[bytes, List[UsmPage]],
+        frame_rate: int = 30,
+        frame_time: int = 0,
+        padding: Union[int, Callable[[int], int]] = 0,
+        channel_number: int = 0,
+        payload_offset: int = 0x18,
+        encoding: str = "UTF-8",
+    ):
+        self.chunk_type = chunk_type
+        self.payload_type = payload_type
+        self.payload = payload
+        self.frame_rate = frame_rate
+        self.frame_time = frame_time
+        self._padding = padding
+        self.channel_number = channel_number
+        self.payload_offset = payload_offset
+        self.encoding = encoding
+
+    @property
+    def padding(self) -> int:
+        """The number of byte padding a chunk will have when packed."""
+        if isinstance(self._padding, int):
+            return self._padding
+
+        if isinstance(self.payload, list):
+            payload_size = len(pack_pages(self.payload, self.encoding))
+        else:
+            payload_size = len(self.payload)
+
+        return self._padding(0x20 + payload_size)
+
+    def __len__(self) -> int:
+        """Returns the packed length of a chunk. Including _padding."""
+        if isinstance(self.payload, list):
+            payload_size = len(pack_pages(self.payload, self.encoding))
+        else:
+            payload_size = len(self.payload)
+
+        if isinstance(self._padding, int):
+            padding = self._padding
+        else:
+            padding = self._padding(0x20 + payload_size)
+
+        return 0x20 + payload_size + padding
+
+    @classmethod
+    def from_bytes(cls, chunk: bytes, encoding: str = "UTF-8") -> UsmChunk:
+        chunk = bytearray(chunk)
+        signature = chunk[:0x4]
+
+        chunksize = int.from_bytes(chunk[0x4:0x8], "big")
+        # r08: 1 byte
+        payload_offset = chunk[0x9]
+        padding = int.from_bytes(chunk[0xA:0xC], "big")
+        channel_number = chunk[0xC]
+        # r0D: 1 byte
+        # r0E: 1 byte
+
+        payload_type = PayloadType.from_int(chunk[0xF] & 0x3)
+
+        frame_time = int.from_bytes(chunk[0x10:0x14], "big")
+        frame_rate = int.from_bytes(chunk[0x14:0x18], "big")
+        # r18: 4 bytes
+        # r1C: 4 bytes
+
+        logging.debug(
+            "UsmChunk: Chunk type: %s, chunk size: %x, r08: %x, payload offset: %x "
+            + "padding: %x, chno: %x, r0D: %x, r0E: %x, payload type: %s "
+            + "frame time: %x, frame rate: %d, r18: %s, r1C: %s",
+            bytes_to_hex(signature),
+            chunksize,
+            chunk[0x8],
+            payload_offset,
+            padding,
+            channel_number,
+            chunk[0xD],
+            chunk[0xE],
+            payload_type,
+            frame_time,
+            frame_rate,
+            bytes_to_hex(chunk[0x18:0x1C]),
+            bytes_to_hex(chunk[0x1C:0x20]),
+        )
+
+        if not is_valid_chunk(signature):
+            raise ValueError(f"Invalid signature: {bytes_to_hex(signature)}")
+
+        payload_begin = 0x08 + payload_offset
+        payload_size = chunksize - padding - payload_offset
+        payload: bytearray = chunk[payload_begin : payload_begin + payload_size]
+
+        # Get pages for header and seek payload types
+        if payload_type in [PayloadType.HEADER, PayloadType.METADATA]:
+            payload: List[UsmPage] = get_pages(payload, encoding)
+            for page in payload:
+                logging.debug("Name: %s, Contents: %s", page.name, page.dict)
+
+        return cls(
+            ChunkType.from_bytes(signature),
+            payload_type,
+            payload,
+            frame_rate,
+            frame_time=frame_time,
+            padding=padding,
+            channel_number=channel_number,
+            payload_offset=payload_begin,
+        )
+
+    def pack(self) -> bytes:
+        result = bytearray()
+        result += self.chunk_type.value
+
+        if isinstance(self.payload, list):
+            payload = pack_pages(self.payload, self.encoding)
+        else:
+            payload = self.payload
+
+        if isinstance(self._padding, int):
+            padding = self._padding
+        else:
+            padding = self._padding(0x20 + len(payload))
+
+        chunksize = 0x18 + len(payload) + padding
+        result += chunksize.to_bytes(4, "big")
+        result += bytes(1)
+        result += (0x18).to_bytes(1, "big")
+        result += padding.to_bytes(2, "big")
+        result += self.channel_number.to_bytes(1, "big")
+        result += bytes(2)
+        result += self.payload_type.value.to_bytes(1, "big")
+        result += self.frame_time.to_bytes(4, "big")
+        result += self.frame_rate.to_bytes(4, "big")
+
+        result += bytearray(8)
+        result += payload
+        result += bytearray(padding)
+        return bytes(result)

wannacri/usm/media/__init__.py

@@ -0,0 +1,7 @@
+from .protocols import UsmVideo, UsmAudio, UsmMedia
+from .video import GenericVideo, Vp9
+from .audio import GenericAudio
+from .tools import (
+    create_video_crid_page,
+    create_video_header_page,
+)

wannacri/usm/media/audio.py

@@ -0,0 +1,26 @@
+from typing import Generator, Optional, List
+
+from .protocols import UsmAudio
+from ..page import UsmPage
+
+
+class GenericAudio(UsmAudio):
+    """Generic audios container used for storing audios
+    channels in Usm files. Use other containers when creating
+    USMs from audios files."""
+
+    def __init__(
+        self,
+        stream: Generator[bytes, None, None],
+        crid_page: UsmPage,
+        header_page: UsmPage,
+        length: int,
+        channel_number: int = 0,
+        metadata_pages: Optional[List[UsmPage]] = None,
+    ):
+        self._stream = stream
+        self._crid_page = crid_page
+        self._header_page = header_page
+        self._length = length
+        self._channel_number = channel_number
+        self._metadata_pages = metadata_pages