PyLibMS 2.0.0__py3-none-any.whl

LMS/Message/MSBTStream.py

@@ -0,0 +1,138 @@
+from typing import BinaryIO
+
+from LMS.Common import LMS_Exceptions
+from LMS.Common.Stream.FileInfo import read_file_info, write_file_info
+from LMS.Common.Stream.Hashtable import read_labels, write_labels
+from LMS.Common.Stream.Section import (read_section_data, write_section,
+                                       write_unsupported_section)
+from LMS.Config.TitleConfig import AttributeConfig, TagConfig
+from LMS.FileIO.Stream import FileReader, FileWriter
+from LMS.Message.MSBT import MSBT
+from LMS.Message.MSBTEntry import MSBTEntry
+from LMS.Message.Section.ATR1 import (read_decoded_atr1, read_encoded_atr1,
+                                      write_decoded_atr1, write_encoded_atr1)
+from LMS.Message.Section.TSY1 import read_tsy1, write_tsy1
+from LMS.Message.Section.TXT2 import read_txt2, write_txt2
+
+
+def read_msbt(
+    stream: BinaryIO,
+    attribute_config: AttributeConfig = None,
+    tag_config: TagConfig = None,
+) -> MSBT:
+    """Reads a MSBT file from a stream.
+
+    :param stream: a stream object.
+
+    ## Usage
+    ```
+    with open(file_path, "rb") as file:
+        msbt = read_msbt(file, "Game.yaml")
+        ...
+    ```"""
+    if stream is None:
+        raise ValueError("Stream must be valid!")
+
+    reader = FileReader(stream)
+    file_info = read_file_info(reader, "MsgStdBn")
+
+    file = MSBT(file_info)
+
+    if attribute_config is not None:
+        file.encoded_attributes = False
+
+    messages = attributes = style_indexes = None
+    for magic, size in read_section_data(reader, file_info.section_count):
+        match magic:
+            case "LBL1":
+                labels, slot_count = read_labels(reader)
+                file.slot_count = slot_count
+            case "ATR1":
+                # Map to shared variable to make assignment cleaner
+                if attribute_config is None:
+                    data = read_encoded_atr1(reader, size)
+                else:
+                    file.encoded_attributes = False
+                    data = read_decoded_atr1(reader, attribute_config)
+
+                attributes, size_per_attribute, string_table = data
+                file.size_per_attribute = size_per_attribute
+                file.attr_string_table = string_table
+            case "TXT2":
+                messages = read_txt2(reader, file_info.encoding, tag_config)
+            case "TSY1":
+                style_indexes = read_tsy1(reader, len(labels))
+            case _:
+                file.unsupported_sections[magic] = reader.read_bytes(size)
+
+        file.section_list.append(magic)
+
+    for i, label in labels.items():
+        text = None if messages is None else messages[i]
+        attribute = None if attributes is None else attributes[i]
+        style_index = None if style_indexes is None else style_indexes[i]
+        file.entries.append(MSBTEntry(label, text, attribute, style_index))
+
+    return file
+
+
+def write_msbt(stream: BinaryIO, file: MSBT) -> None:
+    """Writes a MSBT file to a stream.
+
+    :param stream: the filestream.
+    :param file: the MSBT file object.
+
+    ## Usage
+    ```
+        with open(file_path, "wb") as file:
+            write_msbt(file, msbt)
+            ...
+    ```
+    """
+    if stream is None:
+        raise LMS_Exceptions.LMS_Error("Stream is not valid!")
+
+    if not isinstance(file, MSBT):
+        raise LMS_Exceptions.LMS_Error("File provided is not a MSBT.")
+
+    writer = FileWriter(file.info.encoding)
+    write_file_info(writer, "MsgStdBn", file.info)
+
+    for section in file.section_list:
+        match section:
+            case "LBL1":
+                labels = [entry.name for entry in file.entries]
+                write_section(writer, "LBL1", write_labels, labels, file.slot_count)
+            case "ATR1":
+                attributes = [entry.attribute for entry in file.entries]
+                if file.encoded_attributes:
+                    write_section(
+                        writer,
+                        "ATR1",
+                        write_encoded_atr1,
+                        attributes,
+                        file.size_per_attribute,
+                        file.attr_string_table,
+                    )
+                else:
+                    write_section(
+                        writer,
+                        "ATR1",
+                        write_decoded_atr1,
+                        attributes,
+                        file.size_per_attribute,
+                    )
+            case "TXT2":
+                messages = [entry.message for entry in file.entries]
+                write_section(writer, "TXT2", write_txt2, messages)
+            case "TSY1":
+                style_indexes = [entry.style_index for entry in file.entries]
+                write_section(writer, "TSY1", write_tsy1, style_indexes)
+            case _:
+                write_unsupported_section(
+                    writer, section, file.unsupported_sections[section]
+                )
+
+    writer.seek(0x12)
+    writer.write_uint32(writer.get_stream_size())
+    stream.write(writer.get_data())

LMS/Message/Section/ATR1.py

@@ -0,0 +1,101 @@
+from LMS.Config.Definitions.Attributes import AttributeConfig
+from LMS.Field.LMS_DataType import LMS_DataType
+from LMS.Field.LMS_Field import LMS_DataType, LMS_Field
+from LMS.Field.Stream import read_field, write_field
+from LMS.FileIO.Stream import FileReader, FileWriter
+from LMS.Message.Definitions.LMS_FieldMap import LMS_FieldMap
+
+
+def read_encoded_atr1(
+    reader: FileReader, section_size: int
+) -> tuple[list[bytes], int, list[str]]:
+    absolute_size = section_size + reader.tell()
+
+    attribute_count = reader.read_uint32()
+    size_per_attribute = reader.read_uint32()
+
+    attributes = [reader.read_bytes(size_per_attribute) for _ in range(attribute_count)]
+    string_table = None
+
+    if section_size > 8 + size_per_attribute * attribute_count:
+        string_table = reader.read_bytes(absolute_size - reader.tell())
+
+    return (attributes, size_per_attribute, string_table)
+
+
+def write_encoded_atr1(
+    writer: FileWriter,
+    attributes: list[bytes],
+    size_per_attribute: int,
+    string_table: bytes | None,
+):
+    writer.write_uint32(len(attributes))
+
+    if attributes:
+        writer.write_uint32(size_per_attribute)
+    else:
+        writer.write_uint32(0)
+
+    for attr in attributes:
+        writer.write_bytes(attr)
+
+    if string_table is not None:
+        writer.write_bytes(string_table)
+
+
+def read_decoded_atr1(
+    reader: FileReader, structure: AttributeConfig
+) -> tuple[list[LMS_FieldMap], int, None]:
+    section_start = reader.tell()
+
+    attr_count = reader.read_uint32()
+    size_per_attr = reader.read_uint32()
+
+    attributes = []
+    attr_start = reader.tell()
+
+    for _ in range(attr_count):
+        reader.seek(attr_start)
+
+        attribute = {}
+        for definition in structure.definitions:
+            if definition.datatype is LMS_DataType.STRING:
+                last = reader.tell() + 4
+                reader.seek(section_start + reader.read_uint32())
+                value = LMS_Field(reader.read_str_variable_encoding(), definition)
+
+                reader.seek(last)
+            else:
+                value = read_field(reader, definition)
+
+            attribute[definition.name] = value
+
+        attributes.append(attribute)
+
+        # Move to the start of the next attribute
+        attr_start += size_per_attr
+
+    return attributes, size_per_attr, None
+
+
+def write_decoded_atr1(
+    writer: FileWriter, attributes: list[LMS_FieldMap], size_per_attribute: int
+) -> None:
+    writer.write_uint32(len(attributes))
+    writer.write_uint32(size_per_attribute)
+
+    string_table = []
+    string_offset = 8 + size_per_attribute * len(attributes)
+    for attr in attributes:
+        for field in attr.values():
+            if field.datatype is not LMS_DataType.STRING:
+                write_field(writer, field)
+            else:
+                writer.write_uint32(string_offset)
+                string_offset += (
+                    len(field.value) * writer.encoding.width
+                ) + writer.encoding.width
+                string_table.append(field.value)
+
+    for string in string_table:
+        writer.write_variable_encoding_string(string)

LMS/Message/Section/TSY1.py

@@ -0,0 +1,13 @@
+from LMS.FileIO.Stream import FileReader, FileWriter
+
+
+def read_tsy1(reader: FileReader, message_count: int) -> list[int]:
+    style_indexes = []
+    for _ in range(message_count):
+        style_indexes.append(reader.read_uint32())
+    return style_indexes
+
+
+def write_tsy1(writer: FileWriter, style_indexes: list[int]) -> None:
+    for i in style_indexes:
+        writer.write_uint32(i)

LMS/Message/Section/TXT2.py

@@ -0,0 +1,79 @@
+from LMS.Config.Definitions.Tags import TagConfig
+from LMS.FileIO.Encoding import FileEncoding
+from LMS.FileIO.Stream import FileReader, FileWriter
+from LMS.Message.Definitions.LMS_MessageText import LMS_MessageText
+from LMS.Message.Tag.LMS_Tag import LMS_TagBase
+from LMS.Message.Tag.Stream import read_tag, write_tag
+
+
+def read_txt2(
+    reader: FileReader, encoding: FileEncoding, config: TagConfig
+) -> list[LMS_MessageText]:
+    text_list = []
+    count = reader.read_uint32()
+
+    encoding_format = encoding.to_string_format(reader.big_endian)
+    tag_indicator = b"\x0e" + (b"\x00" * (reader.encoding.width - 1))
+
+    if reader.big_endian:
+        tag_indicator = tag_indicator[::-1]
+
+    for offset in reader.read_offset_array(count):
+        reader.seek(offset)
+
+        text, parts = b"", []
+        while (data := reader.read_bytes(encoding.width)) != encoding.terminator:
+            if data == tag_indicator:
+                # Add all the text before the control tag
+                parts.append(text.decode(encoding_format))
+
+                # Read the tag data
+                group_index = reader.read_uint16()
+                tag_index = reader.read_uint16()
+                param_size = reader.read_uint16()
+
+                tag = read_tag(reader, param_size, group_index, tag_index, config)
+
+                parts.append(tag)
+                text = b""
+            # TODO: Add support for closing tags
+            elif data == b"\x0f":
+                raise NotImplementedError(
+                    "MSBT files with closing tags are unsupported at this time!"
+                )
+            else:
+                text += data
+
+        # Add the remaining text in case there were no control tags
+        if text:
+            parts.append(text.decode(encoding_format))
+
+        message = LMS_MessageText(parts, config)
+        text_list.append(message)
+
+    return text_list
+
+
+def write_txt2(writer: FileWriter, messages: list[LMS_MessageText]) -> None:
+    start = writer.tell()
+    writer.write_uint32(len(messages))
+
+    offset = 4 + 4 * len(messages)
+    for message in messages:
+        writer.write_uint32(offset)
+        next = writer.tell()
+
+        writer.seek(start + offset)
+        text_start = writer.tell()
+        for part in message:
+            if isinstance(part, LMS_TagBase):
+                write_tag(writer, part)
+            else:
+                writer.write_variable_encoding_string(part, False)
+
+        writer.write_bytes(writer.encoding.terminator)
+        offset += writer.tell() - text_start
+        writer.seek(next)
+
+    # Seek to the end of the text to ensure end data can be written properly
+    writer.seek(offset + start)

LMS/Message/Tag/LMS_Tag.py

@@ -0,0 +1,169 @@
+import re
+from abc import ABC, abstractmethod
+
+from LMS.Config.Definitions.Tags import TagConfig
+from LMS.Field.LMS_Field import LMS_Field, cast_value
+from LMS.Message.Definitions.LMS_FieldMap import LMS_FieldMap
+from LMS.Message.Tag.LMS_TagExceptions import LMS_InvalidTagFormatError
+from LMS.Message.Tag.Tag_Formats import (DECODED_FORMAT, ENCODED_FORMAT,
+                                         PARAMETER_FORMAT)
+
+
+class LMS_TagBase(ABC):
+    def __init__(
+        self,
+        group_index: int,
+        tag_index: int,
+        parameters: LMS_FieldMap | list[str] | None = None,
+        group_name: str = None,
+        tag_name: str = None,
+    ):
+        self._group_index = group_index
+        self._tag_index = tag_index
+        self._parameters = parameters or {}
+        self._group_name, self._tag_name = group_name, tag_name
+
+    @property
+    def group_index(self) -> int:
+        """The index of the group."""
+        return self._group_index
+
+    @property
+    def tag_index(self) -> int:
+        """The index of the tag in it's group."""
+        return self._tag_index
+
+    @property
+    def group_name(self) -> str:
+        """The name of the tag group."""
+        return self._group_name
+
+    @property
+    def tag_name(self) -> str:
+        """The name of the tag in the tag group."""
+        return self._tag_name
+
+    @property
+    def parameters(self) -> list[str] | LMS_FieldMap:
+        """The parameters of the tag."""
+        return self._parameters
+
+    @abstractmethod
+    def to_text(self) -> str:
+        """Converts the tag to its string representation."""
+        pass
+
+    @classmethod
+    @abstractmethod
+    def from_string(self, string: str, config: TagConfig | None = None) -> str:
+        pass
+
+
+class LMS_EncodedTag(LMS_TagBase):
+    """A class that represents an encoded tag.
+
+    Example encoded tags:
+        - `[0:3 00-00-00-FF]`
+        - `[0:4]`
+        - `[1:0 01-00-00-CD]`"""
+
+    def __init__(
+        self,
+        group_index: int,
+        tag_index: int,
+        parameters: list[str] = None,
+        group_name: str = None,
+        tag_name: str = None,
+    ):
+        super().__init__(group_index, tag_index, parameters)
+        self._group_name = group_name
+        self._tag_name = tag_name
+
+    def to_text(self) -> str:
+        if self._group_name is not None and self._tag_name is not None:
+            # Determine what to display based on if the names are provided
+            group = self._group_name or self._group_index
+            tag = self._tag_name or self._tag_index
+        else:
+            group, tag = self._group_index, self._tag_index
+
+        if not self._parameters:
+            return f"[{group}:{tag}]"
+
+        parameters = "-".join(self.parameters)
+        return f"[{group}:{tag} {parameters}]"
+
+    @classmethod
+    def from_string(cls, string: str, config: TagConfig | None = None):
+        if match := re.match(ENCODED_FORMAT, string):
+            group, tag = match.group(1), match.group(2)
+            parameters = match.group(3).split("-")
+
+            if group.isdigit() and tag.isdigit():
+                return cls(int(group), int(tag), parameters)
+
+            tag_definition = config.get_definition_by_names(group, tag)
+            return cls(
+                group,
+                tag,
+                parameters,
+                tag_definition.group_name,
+                tag_definition.tag_name,
+            )
+        else:
+            raise LMS_InvalidTagFormatError(
+                f"Invalid encoded tag format detected for tag'{string}'"
+            )
+
+
+class LMS_DecodedTag(LMS_TagBase):
+    """A class that represents a decoded tag.
+
+    Example encoded tags:
+        - `[System:Color r="0" g="255" b="255" a="255"]`
+        - `[System:Pagebreak]`
+        - `[Mii:Nickname buffer="1" type="Text" conversion="None"]`"""
+
+    def __init__(
+        self,
+        group_index: int,
+        tag_index: int,
+        group_name: str,
+        tag_name: str,
+        parameters: LMS_FieldMap = None,
+    ):
+        super().__init__(group_index, tag_index, parameters, group_name, tag_name)
+
+    def to_text(self) -> str:
+        if not self._parameters:
+            return f"[{self.group_name}:{self.tag_name}]"
+
+        parameters = " ".join(
+            f'{key}="{param._value}"' for key, param in self._parameters.items()
+        )
+        return f"[{self.group_name}:{self.tag_name} {parameters}]"
+
+    @classmethod
+    def from_string(cls, tag: str, config: TagConfig):
+        if match := re.match(DECODED_FORMAT, tag):
+            group_name, tag_name = match.group(1), match.group(2)
+            parameters = dict(re.findall(PARAMETER_FORMAT, tag))
+
+            tag_definition = config.get_definition_by_names(group_name, tag_name)
+            for definition in tag_definition.parameters:
+                casted_value = cast_value(
+                    parameters[definition.name], definition.datatype
+                )
+                parameters[definition.name] = LMS_Field(casted_value, definition)
+
+            return cls(
+                tag_definition.group_index,
+                tag_definition.tag_index,
+                group_name,
+                tag_name,
+                parameters,
+            )
+        else:
+            raise LMS_InvalidTagFormatError(
+                f"Invalid decoded tag format detectefd for tag '{tag}'"
+            )

LMS/Message/Tag/LMS_TagExceptions.py

@@ -0,0 +1,7 @@
+class LMS_TagReadingError(Exception): ...
+
+
+class LMS_TagWritingException(Exception): ...
+
+
+class LMS_InvalidTagFormatError(Exception): ...

LMS/Message/Tag/Stream.py

@@ -0,0 +1,149 @@
+from LMS.Config.Definitions.Tags import TagConfig, TagDefinition
+from LMS.Field.LMS_DataType import LMS_DataType
+from LMS.Field.LMS_Field import LMS_Field
+from LMS.Field.Stream import read_field, write_field
+from LMS.FileIO.Stream import FileReader, FileWriter
+from LMS.Message.Definitions.LMS_FieldMap import LMS_FieldMap
+from LMS.Message.Tag.LMS_Tag import LMS_DecodedTag, LMS_EncodedTag, LMS_TagBase
+from LMS.Message.Tag.LMS_TagExceptions import (LMS_TagReadingError,
+                                               LMS_TagWritingException)
+
+
+def read_tag(
+    reader: FileReader,
+    param_size: int,
+    group_index: int,
+    tag_index: int,
+    config: TagConfig | None,
+) -> LMS_EncodedTag | LMS_DecodedTag:
+    end = reader.tell() + param_size
+
+    if config is None:
+        parameters = _read_encoded_parameters(reader, param_size)
+
+        if parameters is None:
+            return LMS_EncodedTag(group_index, tag_index)
+
+        return LMS_EncodedTag(group_index, tag_index, parameters)
+
+    definition = config.get_definition_by_indexes(group_index, tag_index)
+
+    # If the parameters were omitted from the definition but the tag still has defined parameters, read them.
+    # This is to account for encoded tags that group have tag names attatched.
+    if definition.parameters is None and param_size:
+        parameters = _read_encoded_parameters(reader, param_size)
+        return LMS_EncodedTag(
+            group_index,
+            tag_index,
+            parameters,
+            definition.group_name,
+            definition.tag_name,
+        )
+    else:
+        parameters = _read_decoded_parameters(reader, definition)
+
+    reader.seek(end)
+    return LMS_DecodedTag(
+        group_index, tag_index, definition.group_name, definition.tag_name, parameters
+    )
+
+
+def write_tag(writer: FileWriter, tag: LMS_TagBase) -> None:
+    tag_indicator = b"\x0e" + (b"\x00" * (writer.encoding.width - 1))
+    if writer.big_endian:
+        tag_indicator = tag_indicator[::-1]
+
+    writer.write_bytes(tag_indicator)
+    writer.write_uint16(tag.group_index)
+    writer.write_uint16(tag.tag_index)
+
+    if tag.parameters is None:
+        writer.write_uint16(0)
+        return None
+
+    if isinstance(tag, LMS_EncodedTag):
+        _write_encoded_parameters(writer, tag.parameters)
+    else:
+        _write_decoded_parameters(writer, tag)
+
+    return None
+
+
+# --
+def _read_encoded_parameters(reader: FileReader, param_size: int) -> list[str] | None:
+    if param_size == 0:
+        return None
+
+    hex_parameters = reader.read_bytes(param_size).hex().upper()
+    encoded_parameters = [
+        hex_parameters[i : i + 2] for i in range(0, len(hex_parameters), 2)
+    ]
+    return encoded_parameters
+
+
+def _write_encoded_parameters(writer: FileWriter, parameters: list[str] | None) -> None:
+    writer.write_uint16(len(parameters))
+    for param in parameters:
+        writer.write_bytes(bytes.fromhex(param))
+
+
+def _read_decoded_parameters(
+    reader: FileReader, definition: TagDefinition
+) -> LMS_FieldMap:
+    parameters = {}
+    for param in definition.parameters:
+        param_offset = reader.tell()
+        try:
+            if param.datatype is LMS_DataType.STRING:
+                value = LMS_Field(reader.read_len_string_variable_encoding(), param)
+            else:
+                value = read_field(reader, param)
+        # There could be multiple errors related to reading, share the extra info but display the original exception
+        except Exception as e:
+            raise LMS_TagReadingError(
+                f"An error occured reading tag '[{definition.group_name}:{definition.tag_name}]', parameter '{param.name}' at offset {param_offset}"
+            ) from e
+
+        parameters[param.name] = value
+    return parameters
+
+
+def _write_decoded_parameters(writer: FileWriter, tag: LMS_DecodedTag) -> None:
+    param_size = 0
+
+    # Tags are padded by 0xCD if the size is not aligned to the encoding
+    # This can occur before a string parameter, or at the end of the tag.
+    # Set a flag in order to be able to pad the first string correctly
+    needs_padding = False
+    first_string = True
+
+    for field in tag.parameters.values():
+        if field.datatype is LMS_DataType.STRING:
+            param_size += 2 + len(field.value) * writer.encoding.width
+        else:
+            param_size += field.datatype.stream_size
+
+    if param_size % 2 == 1:
+        needs_padding = True
+        param_size += 1
+
+    writer.write_uint16(param_size)
+    for field in tag.parameters.values():
+        try:
+            if field.datatype is LMS_DataType.STRING:
+                if first_string and needs_padding:
+                    writer.write_bytes(b"\xcd")
+                    first_string = False
+                    needs_padding = False
+
+                writer.write_uint16(len(field.value) * writer.encoding.width)
+                writer.write_variable_encoding_string(field.value, False)
+            else:
+                write_field(writer, field)
+        except Exception as e:
+            raise LMS_TagWritingException(
+                f"An error occured writing tag '{tag}', parameter '{field.name}' at offset {writer.tell()}!"
+            ) from e
+
+    if needs_padding:
+        writer.write_bytes(b"\xcd")