psmbuild.py

#!/usr/bin/env python3

"""
PSM Builder form the PSM Toolkit
Copyright (C) 2026 CosmicScale

<https://github.com/CosmicScale/PSBBN-Definitive-English-Patch>

SPDX-License-Identifier: GPL-3.0-or-later

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import sys
import struct
import os

# ------------------------------
# Constants for PSM header
# ------------------------------
PS2S_HEADER_SIZE        = 0x90
OFFSET_MAGIC            = 0x00  # 4 bytes  - "PS2S"
OFFSET_VERSION          = 0x04  # 2 bytes  - Format version
OFFSET_FLAGS            = 0x06  # 2 bytes  - Flags (usually 0)
OFFSET_DATA_OFFSET      = 0x08  # 4 bytes  - Offset to MPEG stream
OFFSET_FIELD_COUNT      = 0x0C  # 4 bytes  - Observed constant (0x04)
OFFSET_INTERNAL_OFFSET  = 0x10  # 4 bytes  - maximum title length (in bytes) that PSBBN will read/display from the title field
OFFSET_TITLE_POINTER    = 0x14  # 4 bytes  - Offset to title (0x50)
OFFSET_THUMB_SIZE       = 0x40  # 4 bytes  - Thumbnail size
OFFSET_THUMB_OFFSET     = 0x44  # 4 bytes  - Thumbnail offset
OFFSET_TITLE            = 0x50  # Title starts here (UTF-8)
TITLE_MAX_LEN           = 0x40  # 64 bytes - Title (UTF-8)

PS2S_VERSION            = 0x0100
PS2S_FLAGS              = 0x0000
PS2S_FIELD_COUNT        = 0x00000004
PS2S_INTERNAL_OFFSET    = 0x00000040
PS2S_TITLE_POINTER      = OFFSET_TITLE

# ==============================
# Utility functions
# ==============================
def align16(value):
    return (value + 15) & ~15

def truncate_utf8_safe(text, max_bytes):
    encoded = bytearray()
    for ch in text:
        ch_bytes = ch.encode("utf-8", errors="ignore")
        if len(encoded) + len(ch_bytes) > max_bytes:
            break
        encoded.extend(ch_bytes)
    return bytes(encoded)

# ==============================
# Build header
# ==============================
def build_psm_header(thumb_size, thumb_offset, data_offset, title_text):
    header = bytearray(PS2S_HEADER_SIZE)

    # Magic & version
    header[OFFSET_MAGIC:OFFSET_MAGIC+4] = b"PS2S"
    struct.pack_into("<H", header, OFFSET_VERSION, PS2S_VERSION)
    struct.pack_into("<H", header, OFFSET_FLAGS, PS2S_FLAGS)

    # Standard fields
    struct.pack_into("<I", header, OFFSET_DATA_OFFSET, data_offset)
    struct.pack_into("<I", header, OFFSET_FIELD_COUNT, PS2S_FIELD_COUNT)
    struct.pack_into("<I", header, OFFSET_INTERNAL_OFFSET, PS2S_INTERNAL_OFFSET)
    struct.pack_into("<I", header, OFFSET_TITLE_POINTER, PS2S_TITLE_POINTER)

    # Thumbnail
    struct.pack_into("<I", header, OFFSET_THUMB_SIZE, thumb_size)
    struct.pack_into("<I", header, OFFSET_THUMB_OFFSET, thumb_offset)

    # Title (UTF-8)
    safe_title = truncate_utf8_safe(title_text, TITLE_MAX_LEN - 1)
    header[OFFSET_TITLE:OFFSET_TITLE + TITLE_MAX_LEN] = b"\x00" * TITLE_MAX_LEN
    header[OFFSET_TITLE:OFFSET_TITLE + len(safe_title)] = safe_title

    return header

# ------------------------------
# Build PSM
# ------------------------------
def build_psm(pss_file, tm2_file, output_file, title_text):

    # Read TM2 texture
    with open(tm2_file, "rb") as f:
        tm2_data = f.read()

    if tm2_data[0:4] != b"TIM2":
        raise ValueError("Input file is not a valid TIM2 texture!")

    # Read MPEG stream
    with open(pss_file, "rb") as f:
        mpeg_data = f.read()

    if mpeg_data[0:4] != b"\x00\x00\x01\xBA":
        raise ValueError("PSS file not valid MPEG-PS!")

    # Offsets
    thumb_offset = PS2S_HEADER_SIZE
    raw_end = thumb_offset + len(tm2_data)
    data_offset = align16(raw_end)
    padding_size = data_offset - raw_end

    # Build header
    header = build_psm_header(
        thumb_size=len(tm2_data),
        thumb_offset=thumb_offset,
        data_offset=data_offset,
        title_text=title_text
    )

    # Write output
    with open(output_file, "wb") as f:
        f.write(header)
        f.write(tm2_data)
        f.write(b"\x00" * padding_size)
        f.write(mpeg_data)

    print(f"[✓] Created PSM: {os.path.basename(output_file)}")

# ------------------------------
# Main
# ------------------------------
if __name__ == "__main__":

    if len(sys.argv) < 3 or len(sys.argv) > 4:
        print("Usage: python3 psmbuild.py <input.pss> <input.tm2> [title]")
        sys.exit(1)

    pss_file = sys.argv[1]
    tm2_file = sys.argv[2]

    # Output file derived from PSS filename
    base_name = os.path.splitext(os.path.basename(pss_file))[0]
    output_file = base_name + ".psm"

    # Optional title argument
    if len(sys.argv) == 4:
        title = sys.argv[3]
    else:
        title = base_name

    build_psm(pss_file, tm2_file, output_file, title_text=title)