Merge pull request #4 from danila-schelkov/decompress-hotfix

fix(decompress): hash reading, comparing file MAGIC before decompression

Author: danila-schelkov

sc_compression/decompressor.py

@@ -1,6 +1,7 @@
 import lzma
 
-from sc_compression.signatures import Signatures, get_signature
+from sc_compression.exceptions import UnknownFileMagicException
+from sc_compression.signatures import Signatures, get_signature, MAGIC_SC, MAGIC_ZSTD, MAGIC_SCLZ
 from sc_compression.utils.reader import Reader
 
 try:
@@ -34,24 +35,29 @@ def decompress(self, buffer: bytes) -> bytes:
         if self.signature == Signatures.SC:
             reader = Reader(buffer, "big")
 
-            reader.read(2)
+            self.check_magic(MAGIC_SC, reader=reader)
+
             self.file_version = reader.read_int32()
             if self.file_version == 4:
                 self.file_version = reader.read_int32()
-            elif self.file_version == 0x05000000:
+
+            if self.file_version == 0x05000000:
                 reader.endian = "little"
 
                 metadata_table_offset = reader.read_int32()  # offset to metadata vtable
                 reader.read(metadata_table_offset)  # metadata
             else:
                 hash_length = reader.read_int32()
                 self.hash = reader.read(hash_length)
+
             return self.decompress(buffer[reader.tell() :])
         elif self.signature == Signatures.SIG:
             return self.decompress(buffer[68:])
         elif self.signature == Signatures.SCLZ and lzham is not None:
             reader = Reader(buffer, "little")
-            reader.read(4)
+
+            self.check_magic(MAGIC_SCLZ, reader=reader)
+
             dict_size_log2 = reader.read_u_int8()
             uncompressed_size = reader.read_int32()
 
@@ -64,9 +70,21 @@ def decompress(self, buffer: bytes) -> bytes:
             compressed = buffer[:5] + b"\xff" * 8 + buffer[9:]
             return decompressor.decompress(compressed)
         elif self.signature == Signatures.ZSTD and zstandard is not None:
-            # the decompress.decompress() API errors with zstd.ZstdError: could not
-            # determine content size in frame header
+            self.check_magic(MAGIC_ZSTD, buffer=buffer)
+
             decompressor = zstandard.ZstdDecompressor()
             return decompressor.decompress(buffer)
         else:
             raise TypeError(self.signature)
+
+    @staticmethod
+    def check_magic(expected_magic: bytes, reader: Reader | None = None, buffer: bytes | None = None) -> None:
+        if reader is not None:
+            magic = reader.read(len(expected_magic))
+        elif buffer is not None:
+            magic = buffer[:len(expected_magic)]
+        else:
+            raise Exception("Cannot find buffer to get magic from.")
+
+        if magic != expected_magic:
+            raise UnknownFileMagicException(magic, magic, expected_magic)

sc_compression/exceptions/__init__.py

@@ -0,0 +1 @@
+from .unknown_file_magic import UnknownFileMagicException

sc_compression/exceptions/unknown_file_magic.py

@@ -0,0 +1,3 @@
+class UnknownFileMagicException(Exception):
+    def __init__(self, magic: bytes, expected_magic: bytes, *args):
+        super().__init__(f"Unknown file magic: {magic.hex()}, while expected was {expected_magic.hex()}", *args)

sc_compression/signatures.py

@@ -1,6 +1,11 @@
 import re
 from enum import Enum
 
+MAGIC_SC = b"SC"
+MAGIC_SCLZ = b"SCLZ"
+MAGIC_SIG = b"Sig:"
+MAGIC_ZSTD = b"\x28\xb5\x2f\xfd"
+
 
 class Signatures(Enum):
     NONE = 0
@@ -16,14 +21,14 @@ def get_signature(buffer: bytes) -> Signatures:
 
     if re.match(b"\x00\x00?\x00", buffer[1:5]):
         signature = Signatures.LZMA
-    elif buffer.startswith(b"\x28\xb5\x2f\xfd"):
+    elif buffer.startswith(MAGIC_ZSTD):
         signature = Signatures.ZSTD
 
-    if buffer.startswith(b"SCLZ"):
+    if buffer.startswith(MAGIC_SCLZ):
         signature = Signatures.SCLZ
-    elif buffer.startswith(b"SC"):
+    elif buffer.startswith(MAGIC_SC):
         signature = Signatures.SC
-    elif buffer.startswith(b"Sig:"):
+    elif buffer.startswith(MAGIC_SIG):
         signature = Signatures.SIG
 
     return signature

setup.py

@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="sc-compression",
-    version="0.6.4",
+    version="0.6.5",
     author="Vorono4ka",
     author_email="[email protected]",
     description="SC Compression",