Fix resources extracted from pak_wii being unusable

src/retro_data_structures/compression.py

@@ -97,4 +97,29 @@
  ]
 
 
+class LZOCompressedBlockCorruption(LZOCompressedBlock):
+ # Distinct implementation for Corruption's indexing needs :
+ # Corruption needs distinct indices for the current segment and the current block within the segment
+ def _actual_segment_size(self, context):
+ mem = context._index
+ context._index = context._._index
+
+ decompressed_size = construct.evaluate(self.decompressed_size, context)
+ segment_size = construct.evaluate(self.segment_size, context)
+ context._index = mem
+
+ previous_segments = context._index * segment_size
+ if previous_segments > decompressed_size:
+ # This segment is redundant!
+ raise construct.StopFieldError
+
+ elif previous_segments + segment_size > decompressed_size:
+ # Last segment
+ return decompressed_size - previous_segments
+
+ else:
+ # Another segment with this size
+ return segment_size
+
+
 ZlibCompressedBlock = construct.Compressed(construct.GreedyBytes, "zlib", level=9)

src/retro_data_structures/formats/pak_wii.py

@@ -4,12 +4,12 @@
 from typing import TYPE_CHECKING
 
 import construct
-from construct import Bytes, Const, FocusedSeq, IfThenElse, Int32ub, PrefixedArray, Rebuild, Struct, len_, this
+from construct import Bytes, Const, IfThenElse, Int32ub, PrefixedArray, Struct
 
 from retro_data_structures import game_check
 from retro_data_structures.base_resource import AssetId, AssetType, Dependency
 from retro_data_structures.common_types import AssetId64, FourCC, String
-from retro_data_structures.compression import LZOCompressedBlock, ZlibCompressedBlock
+from retro_data_structures.compression import LZOCompressedBlockCorruption
 from retro_data_structures.construct_extensions.alignment import AlignTo
 from retro_data_structures.construct_extensions.dict import make_dict
 
@@ -70,14 +70,70 @@
  _resources_end=construct.Tell,
 )
 
-CompressedPakResource = FocusedSeq(
- "data",
- decompressed_size=Rebuild(Int32ub, len_(this.data)),
- # Added Zlib check for DKCR
- data=IfThenElse(game_check.uses_lzo, LZOCompressedBlock(this.decompressed_size), ZlibCompressedBlock),
+CMPD_Pak_Resource = Struct(
+ magic=Const(b"CMPD"),
+ block_count=Int32ub,
+ block_header=construct.Array(
+ construct.this.block_count,
+ Struct(
+ flag=construct.Byte,
+ compressed_size=construct.Int24ub,
+ uncompressed_size=construct.Int32ub,
+ ),
+ ),
+ blocks=construct.Array(
+ construct.this.block_count,
+ IfThenElse(
+ lambda this: this.block_header[this._index].compressed_size
+ < this.block_header[this._index].uncompressed_size,
+ Struct(
+ block=LZOCompressedBlockCorruption(lambda this: this._.block_header[this._index].uncompressed_size),
+ ),
+ Bytes(lambda this: this.block_header[this._index].uncompressed_size),
+ ),
+ ),
 )
 
 
+class CMPD_Pak_Adapter(construct.Adapter):
+ def _decode(self, obj, context, path):
+ return b"".join([block if type(block) is bytes else block.block for block in obj.blocks])
+
+ # Going to rip a page out of PWE's book and compress everything in a single block
+ def _encode(self, uncompressed, context, path):
+ res = construct.Container(
+ [
+ ("magic", b"CMPD"),
+ ("block_count", 1),
+ (
+ "block_header",
+ construct.ListContainer(
+ [
+ construct.Container(
+ [
+ ("flag", 0xA0),
+ ("compressed_size", None),
+ ("uncompressed_size", len(uncompressed)),
+ ]
+ ),
+ ]
+ ),
+ ),
+ (
+ "blocks",
+ construct.ListContainer(
+ LZOCompressedBlockCorruption(len(uncompressed))._encode(uncompressed, context, path)
+ ),
+ ),
+ ]
+ )
+ res.block_header[0].compressed_size = len(res.blocks[0])
+ return res
+
+
+CompressedPakResource = CMPD_Pak_Adapter(CMPD_Pak_Resource)
+
+
 @dataclasses.dataclass
 class PakFile:
  asset_id: AssetId
@@ -89,12 +145,12 @@
 
  def get_decompressed(self, target_game: Game) -> bytes:
  if self.uncompressed_data is None:
- self.uncompressed_data = CompressedPakResource.parse(self.compressed_data, target_game=target_game)
+ self.uncompressed_data = CMPD_Pak_Resource.parse(self.compressed_data, target_game=target_game)
  return self.uncompressed_data
 
  def get_compressed(self, target_game: Game) -> bytes:
  if self.compressed_data is None:
- self.compressed_data = CompressedPakResource.build(self.uncompressed_data, target_game=target_game)
+ self.compressed_data = CMPD_Pak_Resource.build(self.uncompressed_data, target_game=target_game)
  return self.compressed_data
 
  def set_new_data(self, data: bytes):

tests/formats/test_pak_wii.py

@@ -1,25 +1,14 @@
 from __future__ import annotations
 
 # The two following imports are only used by file tests
-# from glob import glob
-# from os import path
+import pytest
+
 from retro_data_structures.formats.pak import Pak
-from retro_data_structures.formats.pak_wii import (
- PAK_WII,
- PAKNoData,
-)
-
-# The two following classes are only used by file tests
-# PakFile,
-# PakBody
+from retro_data_structures.formats.pak_wii import PAK_WII, CompressedPakResource, PAKNoData
 from retro_data_structures.game_check import Game
 
 # ruff: noqa: E501
 
-# The following variables are only used for the file tests and should be set before executing said tests locally
-# pak_target = "."
-# pak_build_dir = "."
-
 paks = {
  "FrontEnd",
  "GuiDVD",
@@ -44,6 +33,48 @@
 }
 
 
+@pytest.fixture(name="compressed_resources")
+def _compressed_resources():
+ """
+ The resources can be found in Metroid3.pak
+ """
+ return [
+ { # 2 segments resource
+ "compressed": 1,
+ "asset": {"type": "TXTR", "id": 4849971089334802081},
+ "contents": {
+ "data": b"CMPD\x00\x00\x00\x02\x00\x00\x00\x0c\x00\x00\x00\x0c\xc0\x00\x00\x11\x00\x00\x00\xa0"
+ b"\x00\x00\x00\n\x00\x10\x00\x10\x00\x00\x00\x02\x00\x0f\x16\x01\xe0\x02\xa0\xaa@\x00 "
+ b"w\x1c\x00\x11\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff",
+ "value": (
+ b"\x00\x00\x00\n\x00\x10\x00\x10\x00\x00\x00\x02\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02"
+ b"\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01"
+ b"\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa"
+ b"\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa"
+ b"\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02"
+ b"\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01"
+ b"\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa"
+ b"\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa\x01\xe0\x02\xa0\xaa\xaa\xaa\xaa"
+ ),
+ },
+ },
+ { # 1 segment resource
+ "compressed": 1,
+ "asset": {"type": "CSKR", "id": 2135532429836327754},
+ "contents": {
+ "data": b"CMPD\x00\x00\x00\x01\xa0\x00\x00*\x00\x00\x00J\x00(\x19SKIN\x00\x00\x00\x02M\x00\x01"
+ b"\xcf\x00\x03?\x80W\x00\x00\n\x1as\x00\x00\x00\xff/\x00\x00\xdc\x04+\x02\x00\x01\x00\x11\x00\x00"
+ b"\xff\xff\xff\xff\xff\xff",
+ "value": (
+ b"SKIN\x00\x00\x00\x02\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x03?\x80\x00\x00\x00\x00\n"
+ b"\x1a\x00\x00\x00\n\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"
+ b"\xff\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00"
+ ),
+ },
+ },
+ ]
+
+
 def test_identical_when_keep_data(prime3_iso_provider):
  game = Game.CORRUPTION
  for pakfile in paks:
@@ -77,65 +108,17 @@ def test_compare_header_keep_data(prime3_iso_provider):
  assert custom_sizes == raw_sizes
 
 
-# The following tests are what I call file tests :
-# They produce or read local files specified by the two global variables pak_target and pak_build_dir
-# They are NOT to be executed as tests by CI and are only here for reviewing the testing process
-
-# def test_write_new_pak():
-# game = Game.CORRUPTION
-# files = [
-# PakFile(0xDEADBEEF, "STRG", False, b"abcdefg", None),
-# PakFile(0xDEADD00D, "STRG", False, b"hijklmn", None)
-# ]
-# body = PakBody(b"joe mama so fat ", [
-# ("Hey its me Jack Block from minecraft",
-# Dependency("STRG", 0xDEADBEEF))
-# ],
-# files
-# )
-
-# output_pak = Pak(body, game)
-# encoded = output_pak.build()
-
-# with open(pak_target, "wb") as fd :
-# fd.write(encoded)
-
-# def test_build_from_extracted_pak():
-# game = Game.CORRUPTION
-
-# files = []
-# for file in glob(pak_build_dir + "/*"):
-# asset_id, asset_type = file.split(".")
-# asset_id = int(path.basename(asset_id), 16)
-
-# data = b""
-# with open(file, "rb") as fd:
-# data = fd.read()
-
-# files.append(PakFile(asset_id, asset_type, False, data, None))
-
-# body = PakBody(b"\x1B\x62\xF7\xCA\x15\x60\xB1\x85\xC1\xE1\x09\x43\x99\x4F\xB9\xAC", [
-# ("03b_Bryyo_Fire_#SERIAL#",
-# Dependency("MLVL", 0x9BA9292D588D6EB8)),
-# ("03b_Bryyo_Reptilicus_#SERIAL#",
-# Dependency("MLVL", 0x9F059B53561A9695)),
-# ("03b_Bryyo_Ice_#SERIAL#",
-# Dependency("MLVL", 0xB0D67636D61F3868))
-# ],
-# files
-# )
-
-# output_pak = Pak(body, game)
-# encoded = output_pak.build()
-
-# with open(pak_target, "wb") as fd:
-# fd.write(encoded)
-
-# def test_parse_new_pak():
-# game = Game.CORRUPTION
-
-# with open(pak_target, "rb") as fd:
-# raw = fd.read()
-
-# decoded = Pak.parse(raw, game)
-# return decoded
+def test_corruption_resource_decode(compressed_resources):
+ for compressed_resource in compressed_resources:
+ decoded = CompressedPakResource.parse(compressed_resource["contents"]["data"], target_game=Game.CORRUPTION)
+
+ assert len(decoded) == len(compressed_resource["contents"]["value"])
+ assert decoded == compressed_resource["contents"]["value"]
+
+
+def test_corruption_resource_encode_decode(compressed_resources):
+ for compressed_resource in compressed_resources:
+ raw = compressed_resource["contents"]["value"]
+ decoded = CompressedPakResource.build(raw, target_game=Game.CORRUPTION)
+ encoded = CompressedPakResource.parse(decoded, target_game=Game.CORRUPTION)
+ assert raw == encoded