-
Notifications
You must be signed in to change notification settings - Fork 8
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support reading the headers for Corruption paks
- Loading branch information
1 parent
3479a71
commit fb2720b
Showing
3 changed files
with
136 additions
and
33 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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, | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
), | ||
) |