Support reading the headers for Corruption paks

src/retro_data_structures/construct_extensions/dict.py

@@ -1,38 +1,72 @@
 from __future__ import annotations
 
-from construct import Adapter, PrefixedArray, Struct, VarInt
+import construct
 
 from retro_data_structures.common_types import String
 
 
-class DictAdapter(Adapter):
- def __init__(self, subcon, objisdict=True):
- if not objisdict:
- subcon = Struct("*Key" / VarInt, "Value" / subcon)
- super().__init__(PrefixedArray(VarInt, subcon))
- self.objisdict = objisdict
-
- def _decode(self, obj, context, path):
- D = {}
- for v in obj:
- if self.objisdict:
- D[v["*Key"]] = v
- del v["*Key"]
- else:
- D[v["*Key"]] = v["Value"]
- return D
-
- def _encode(self, obj, context, path):
- L = []
- for k, v in obj.items():
- new_item = v
- if self.objisdict:
- v["*Key"] = k
- else:
- new_item = {"*Key": k, "Value": v}
- L.append(new_item)
- return L
-
-
-def DictStruct(*fields):
- return Struct(*fields, "*Key" / String)
+class DictConstruct(construct.Construct):
+ def __init__(
+ self,
+ key_type: construct.Construct,
+ value_type: construct.Construct,
+ count_type: construct.Construct = construct.Int32ub,
+ ):
+ super().__init__()
+ self.key_type = key_type
+ self.value_type = value_type
+ self.count_type = count_type
+
+ assert not self.key_type.flagbuildnone
+
+ def _parse(self, stream, context, path) -> construct.Container:
+ field_count = self.count_type._parsereport(stream, context, path)
+
+ result = construct.Container()
+
+ for i in range(field_count):
+ field_path = f"{path}.field_{i}"
+ key = self.key_type._parsereport(stream, context, field_path)
+ value = self.value_type._parsereport(stream, context, field_path)
+ result[key] = value
+
+ return result
+
+ def _build(self, obj: construct.Container, stream, context, path):
+ self.count_type._build(len(obj), stream, context, path)
+
+ for i, (key, value) in enumerate(obj.items()):
+ field_path = f"{path}.field_{i}"
+ self.key_type._build(key, stream, context, field_path)
+ self.value_type._build(value, stream, context, field_path)
+
+ def _emitparse(self, code):
+ return "Container(({}, {}) for i in range({}))".format(
+ self.key_type._compileparse(code),
+ self.value_type._compileparse(code),
+ self.count_type._compileparse(code),
+ )
+
+ def _emitbuild(self, code):
+ fname = f"build_dict_{code.allocateId()}"
+ block = f"""
+ def {fname}(key, value, io, this):
+ obj = key
+ {self.key_type._compilebuild(code)}
+
+ obj = value
+ {self.value_type._compilebuild(code)}
+ """
+ code.append(block)
+ return (
+ f"(reuse(len(obj), "
+ f"lambda obj: {self.count_type._compilebuild(code)}), "
+ f"list({fname}(key, value, io, this) for key, value in obj.items()), obj)[2]"
+ )
+
+
+def make_dict(value: construct.Construct, key=String):
+ return DictConstruct(
+ key_type=key,
+ value_type=value,
+ )

src/retro_data_structures/formats/pak.py

@@ -3,14 +3,16 @@
 import typing
 
 from retro_data_structures.base_resource import AssetId, RawResource
-from retro_data_structures.formats import pak_gc, pak_wiiu
+from retro_data_structures.formats import pak_gc, pak_wii, pak_wiiu
 from retro_data_structures.formats.pak_gc import PakBody, PakFile
 from retro_data_structures.game_check import Game
 
 
 def _pak_for_game(game: Game):
  if game == Game.PRIME_REMASTER:
  return pak_wiiu.PAK_WIIU
+ elif game >= Game.CORRUPTION:
+ raise ValueError("Unsupported game")
  else:
  return pak_gc.PAK_GC
 
@@ -27,6 +29,8 @@ def __init__(self, raw: PakBody, target_game: Game):
  def header_for_game(game: Game):
  if game == Game.PRIME_REMASTER:
  return pak_wiiu.PakWiiUNoData
+ elif game >= Game.CORRUPTION:
+ return pak_wii.PAKNoData
  else:
  return pak_gc.PAKNoData
 

src/retro_data_structures/formats/pak_wii.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import construct
+from construct import Bytes, Const, Int32ub, PrefixedArray, Struct
+
+from retro_data_structures.common_types import AssetId64, FourCC, String
+from retro_data_structures.construct_extensions.dict import make_dict
+
+PAKHeader = construct.Aligned(
+ 64,
+ Struct(
+ version=Const(2, Int32ub),
+ header_size=Int32ub,
+ md5_hash=Bytes(16),
+ ),
+)
+
+ConstructResourceHeader = Struct(
+ compressed=Int32ub,
+ asset_type=FourCC,
+ asset_id=AssetId64,
+ size=Int32ub,
+ offset=Int32ub,
+)
+
+
+def _emitparse_header(code: construct.CodeGen) -> str:
+ code.append("ResourceHeader_Format = struct.Struct('>LLQLL')")
+ code.append(
+ """def _create_resource_header(compressed, asset_type, asset_id, size, offset):
+ return Container(compressed=compressed, asset_type=asset_type.to_bytes(4, "big").decode("ascii"),
+ asset_id=asset_id, size=size, offset=offset)
+ """
+ )
+ return "_create_resource_header(*ResourceHeader_Format.unpack(io.read(24)))"
+
+
+ConstructResourceHeader._emitparse = _emitparse_header
+
+PAKNoData = Struct(
+ _start=construct.Tell,
+ _header=PAKHeader,
+ table_of_contents=construct.Aligned(64, make_dict(Int32ub, FourCC)),
+ _named_resources_start=construct.Tell,
+ named_resources=construct.Aligned(
+ 64,
+ PrefixedArray(
+ Int32ub,
+ Struct(
+ name=String,
+ asset_type=FourCC,
+ asset_id=AssetId64,
+ ),
+ ),
+ ),
+ _resources_start=construct.Tell,
+ _resources_start_assert=construct.Check(
+ construct.this.table_of_contents.STRG == construct.this._resources_start - construct.this._named_resources_start
+ ),
+ resources=construct.Aligned(64, PrefixedArray(Int32ub, ConstructResourceHeader)),
+ _resources_end=construct.Tell,
+ _resources_end_assert=construct.Check(
+ construct.this.table_of_contents.RSHD == construct.this._resources_end - construct.this._resources_start
+ ),
+)