feat(handler): add flush method to decompression objects

Author: rxpha3l

python/unblob/handlers/compression/uzip.py

@@ -1,10 +1,10 @@
+import lzma
 import re
 import zlib
-from lzma import LZMADecompressor
 from pathlib import Path
-from typing import Optional
+from typing import Callable, Optional
 
-from pyzstd import EndlessZstdDecompressor
+import pyzstd
 
 from unblob.file_utils import (
     Endian,
@@ -39,10 +39,39 @@
 LZMA_COMPRESSION = "#!/bin/sh\x0a#L3.0\x0a"
 ZSTD_COMPRESSION = "#!/bin/sh\x0a#Z4.0\x20"
 
-DECOMPRESS_METHOD = {
-    ZLIB_COMPRESSION.encode(): zlib.decompressobj,
+
+class Decompressor:
+    DECOMPRESSOR: Callable
+
+    def __init__(self):
+        self._decompressor = self.DECOMPRESSOR()
+
+    def decompress(self, data: bytes) -> bytes:
+        return self._decompressor.decompress(data)
+
+    def flush(self) -> bytes:
+        return b""
+
+
+class LZMADecompressor(Decompressor):
+    DECOMPRESSOR = lzma.LZMADecompressor
+
+
+class ZLIBDecompressor(Decompressor):
+    DECOMPRESSOR = zlib.decompressobj
+
+    def flush(self) -> bytes:
+        return self._decompressor.flush()
+
+
+class ZSTDDecompressor(Decompressor):
+    DECOMPRESSOR = pyzstd.EndlessZstdDecompressor
+
+
+DECOMPRESS_METHOD: dict[bytes, type[Decompressor]] = {
+    ZLIB_COMPRESSION.encode(): ZLIBDecompressor,
     LZMA_COMPRESSION.encode(): LZMADecompressor,
-    ZSTD_COMPRESSION.encode(): EndlessZstdDecompressor,
+    ZSTD_COMPRESSION.encode(): ZSTDDecompressor,
 }
 
 
@@ -67,6 +96,7 @@ def extract(self, inpath: Path, outdir: Path):
                 decompressor = decompressor_cls()
                 for chunk in iterate_file(infile, current_offset, compressed_len):
                     outfile.write(decompressor.decompress(chunk))
+                outfile.write(decompressor.flush())
         return ExtractResult(reports=fs.problems)