From 7c8230e630352e37bbf5a6639b2ad0c9c7ed62c9 Mon Sep 17 00:00:00 2001
From: Yay5379 <61430523+Yay5379@users.noreply.github.com>
Date: Sat, 30 Mar 2024 13:17:55 -0500
Subject: [PATCH] fixed bug with zlib decompression

---
 .../{parse_replay.py => parse_unit_list.py}   |  4 +--
 .../formats/wrpl_parser.py                    | 10 ++----
 src/wt_client_replay_parser/wrpl_unpacker.py  | 31 +++++++------------
 3 files changed, 16 insertions(+), 29 deletions(-)
 rename src/wt_client_replay_parser/formats/{parse_replay.py => parse_unit_list.py} (95%)

diff --git a/src/wt_client_replay_parser/formats/parse_replay.py b/src/wt_client_replay_parser/formats/parse_unit_list.py
similarity index 95%
rename from src/wt_client_replay_parser/formats/parse_replay.py
rename to src/wt_client_replay_parser/formats/parse_unit_list.py
index cc97863..ab4cc6f 100644
--- a/src/wt_client_replay_parser/formats/parse_replay.py
+++ b/src/wt_client_replay_parser/formats/parse_unit_list.py
@@ -3,7 +3,7 @@
 import sys
 import json
 
-def parse_replay(file):
+def parse_replay_unit_list(file):
 
     units = []
     print(f"parsing {file}")
@@ -112,7 +112,7 @@ def main():
 
     print(f"parsing replay in {file}")
 
-    data = parse_replay(file)
+    data = parse_replay_unit_list(file)
 
     folder_name = os.path.basename(file)
     file_path = os.getcwd()
diff --git a/src/wt_client_replay_parser/formats/wrpl_parser.py b/src/wt_client_replay_parser/formats/wrpl_parser.py
index e69f34b..cd3be79 100644
--- a/src/wt_client_replay_parser/formats/wrpl_parser.py
+++ b/src/wt_client_replay_parser/formats/wrpl_parser.py
@@ -67,17 +67,11 @@ def FatBlockStream(sz: t.Union[int, callable, None] = None) -> ct.Construct:
     con = ct.GreedyBytes if sz is None else ct.Bytes(sz)
     return ct.RestreamData(con, bin.Fat)
 
-
-def ZlibStream(sz: t.Union[int, callable, None] = None):
-    con = ct.GreedyBytes if sz is None else ct.Bytes(sz)
-    return ct.RestreamData(con, ct.Compressed(con, 'zlib'))
-
-
 WRPLCliFile = ct.Struct(
     'header' / Header,
     ct.Bytes(2),
     'm_set' / FatBlockStream(this.header.m_set_size),
     'wrplu_offset' / ct.Tell,
-    'wrplu' / ZlibStream(this.header.rez_offset - this.wrplu_offset),
+    'wrplu' / ct.Bytes(this.header.rez_offset - this.wrplu_offset),
     'rez' / bin.Fat,
-)
+)
\ No newline at end of file
diff --git a/src/wt_client_replay_parser/wrpl_unpacker.py b/src/wt_client_replay_parser/wrpl_unpacker.py
index 7a38ba5..692d544 100644
--- a/src/wt_client_replay_parser/wrpl_unpacker.py
+++ b/src/wt_client_replay_parser/wrpl_unpacker.py
@@ -1,24 +1,16 @@
 import argparse
-import json
 from pathlib import Path
+import json
+import zlib
 import typing as t
 import os
 import sys
 import construct as ct
 from blk.types import Section
 import blk.text as txt
-try:
-    from formats.wrpl_parser import WRPLCliFile
-except ImportError:
-    from formats.wrpl_parser import WRPLCliFile
-try:
-    from formats.parse_datablocks import parse_datablocks
-except ImportError:
-    from formats.parse_datablocks import parse_datablocks
-try:
-    from formats.parse_replay import parse_replay
-except ImportError:
-    from formats.parse_replay import parse_replay
+from formats.wrpl_parser import WRPLCliFile
+from formats.parse_datablocks import parse_datablocks
+from formats.parse_unit_list import parse_replay_unit_list
 
 def serialize_text(root: Section, ostream: t.TextIO):
     txt.serialize(root, ostream, dialect=txt.StrictDialect)
@@ -57,9 +49,14 @@ def main():
             serialize_text(section, ostream)
 
     out_path = out_dir / 'wrplu.bin'
-    out_path.write_bytes(parsed.wrplu)
+    out_path.write_bytes(zlib.decompress(parsed.wrplu))
     parse_datablocks(out_path)
 
+    data = parse_replay_unit_list(out_path)
+
+    with create_text(f'{out_dir}/units.json') as ostream:
+        json.dump(data, ostream, indent=2)
+
     out_path = out_dir / 'info.blk'
     info=(
         f'wrpl_version:i={parsed.header.version}\n'
@@ -81,14 +78,10 @@ def main():
     with create_text(out_path) as ostream:
         print(info, file=ostream)
     
-    data = parse_replay(f'{out_dir}/wrplu.bin')
-    with create_text(f'{out_dir}/units.json') as ostream:
-        json.dump(data, ostream, indent=2)
-
     print(f'{replay.name} => {out_dir}')
 
     return 0
 
 
 if __name__ == '__main__':
-    sys.exit(main())
+    sys.exit(main())
\ No newline at end of file