Merge pull request #395 from onekey-sec/feat-zstd-format

martonilles · web-flow · commit 3bd61b1e99cc · 2022-06-14T14:12:07.000+02:00
Add ZSTD compression format
diff --git a/.github/actions/setup-dependencies/action.yml b/.github/actions/setup-dependencies/action.yml
@@ -7,7 +7,7 @@ runs:
   using: "composite"
   steps:
     - name: Install 3rd party from apt
-      run: sudo apt install e2fsprogs p7zip-full unar zlib1g-dev liblzo2-dev lzop lziprecover img2simg libhyperscan-dev
+      run: sudo apt install e2fsprogs p7zip-full unar zlib1g-dev liblzo2-dev lzop lziprecover img2simg libhyperscan-dev zstd
       shell: bash
 
     - name: Install sasquatch
diff --git a/Dockerfile b/Dockerfile
@@ -21,7 +21,8 @@ RUN apt-get update && apt-get install --no-install-recommends -y \
     xz-utils \
     zlib1g-dev \
     libmagic1 \
-    libhyperscan5
+    libhyperscan5 \
+    zstd
 RUN curl -L -o sasquatch_1.0_amd64.deb https://github.com/onekey-sec/sasquatch/releases/download/sasquatch-v1.0/sasquatch_1.0_amd64.deb \
     && dpkg -i sasquatch_1.0_amd64.deb \
     && rm -f sasquatch_1.0_amd64.deb
diff --git a/default.nix b/default.nix
@@ -16,6 +16,7 @@
 , unar
 , file
 , hyperscan
+, zstd
 }:
 
 let
@@ -30,6 +31,7 @@ let
     sasquatch.bigEndian
     simg2img
     unar
+    zstd
   ];
 
   self = mkPoetryApp {
diff --git a/tests/integration/compression/zstd/__input__/lorem.both.zstd b/tests/integration/compression/zstd/__input__/lorem.both.zstd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81468a5c3508970bed02d0824e394497407d0964bc37556de141347f1e320317
+size 2632
diff --git a/tests/integration/compression/zstd/__input__/lorem.prefixed.zstd b/tests/integration/compression/zstd/__input__/lorem.prefixed.zstd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c79b72ddd77445e2ba70810ae04eef63e8dfb0dce9638200e418ea4a670b9f5c
+size 2628
diff --git a/tests/integration/compression/zstd/__input__/lorem.suffixed.zstd b/tests/integration/compression/zstd/__input__/lorem.suffixed.zstd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37aab92d3ca37b75e87a30e7656d26ecd580b88a270cd09df8bcf7a0edd11177
+size 2628
diff --git a/tests/integration/compression/zstd/__input__/lorem.truncated.zstd b/tests/integration/compression/zstd/__input__/lorem.truncated.zstd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2619fb8093ec1eaab0669e4622a5ffe7bce9c28d03b74858a42b9c73c02b0b8
+size 1312
diff --git a/tests/integration/compression/zstd/__input__/lorem.zstd b/tests/integration/compression/zstd/__input__/lorem.zstd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52738e82d35ccdb2c85e5dfcb30f8c320a7ae7e1b427203cd9e566f3d8b97820
+size 2624
diff --git a/tests/integration/compression/zstd/__output__/lorem.both.zstd_extract/0-4.unknown b/tests/integration/compression/zstd/__output__/lorem.both.zstd_extract/0-4.unknown
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36a4aba0dcec23e926f0692065524808479c33f4501bae79348f8f005f10b21c
+size 4
diff --git a/tests/integration/compression/zstd/__output__/lorem.both.zstd_extract/2628-2632.unknown b/tests/integration/compression/zstd/__output__/lorem.both.zstd_extract/2628-2632.unknown
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36a4aba0dcec23e926f0692065524808479c33f4501bae79348f8f005f10b21c
+size 4
diff --git a/tests/integration/compression/zstd/__output__/lorem.both.zstd_extract/4-2628.zstd b/tests/integration/compression/zstd/__output__/lorem.both.zstd_extract/4-2628.zstd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52738e82d35ccdb2c85e5dfcb30f8c320a7ae7e1b427203cd9e566f3d8b97820
+size 2624
diff --git a/tests/integration/compression/zstd/__output__/lorem.both.zstd_extract/4-2628.zstd_extract/4-2628 b/tests/integration/compression/zstd/__output__/lorem.both.zstd_extract/4-2628.zstd_extract/4-2628
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb3492cf53eff7ec69d48b65b923668dd36e80524a2c45d52eae45597a8c8502
+size 6662
diff --git a/tests/integration/compression/zstd/__output__/lorem.prefixed.zstd_extract/0-4.unknown b/tests/integration/compression/zstd/__output__/lorem.prefixed.zstd_extract/0-4.unknown
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36a4aba0dcec23e926f0692065524808479c33f4501bae79348f8f005f10b21c
+size 4
diff --git a/tests/integration/compression/zstd/__output__/lorem.prefixed.zstd_extract/4-2628.zstd b/tests/integration/compression/zstd/__output__/lorem.prefixed.zstd_extract/4-2628.zstd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52738e82d35ccdb2c85e5dfcb30f8c320a7ae7e1b427203cd9e566f3d8b97820
+size 2624
diff --git a/tests/integration/compression/zstd/__output__/lorem.prefixed.zstd_extract/4-2628.zstd_extract/4-2628 b/tests/integration/compression/zstd/__output__/lorem.prefixed.zstd_extract/4-2628.zstd_extract/4-2628
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb3492cf53eff7ec69d48b65b923668dd36e80524a2c45d52eae45597a8c8502
+size 6662
diff --git a/tests/integration/compression/zstd/__output__/lorem.suffixed.zstd_extract/0-2624.zstd b/tests/integration/compression/zstd/__output__/lorem.suffixed.zstd_extract/0-2624.zstd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52738e82d35ccdb2c85e5dfcb30f8c320a7ae7e1b427203cd9e566f3d8b97820
+size 2624
diff --git a/tests/integration/compression/zstd/__output__/lorem.suffixed.zstd_extract/0-2624.zstd_extract/0-2624 b/tests/integration/compression/zstd/__output__/lorem.suffixed.zstd_extract/0-2624.zstd_extract/0-2624
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb3492cf53eff7ec69d48b65b923668dd36e80524a2c45d52eae45597a8c8502
+size 6662
diff --git a/tests/integration/compression/zstd/__output__/lorem.suffixed.zstd_extract/2624-2628.unknown b/tests/integration/compression/zstd/__output__/lorem.suffixed.zstd_extract/2624-2628.unknown
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36a4aba0dcec23e926f0692065524808479c33f4501bae79348f8f005f10b21c
+size 4
diff --git a/tests/integration/compression/zstd/__output__/lorem.truncated.zstd_extract/.gitkeep b/tests/integration/compression/zstd/__output__/lorem.truncated.zstd_extract/.gitkeep
diff --git a/tests/integration/compression/zstd/__output__/lorem.zstd_extract/lorem b/tests/integration/compression/zstd/__output__/lorem.zstd_extract/lorem
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb3492cf53eff7ec69d48b65b923668dd36e80524a2c45d52eae45597a8c8502
+size 6662
diff --git a/unblob/handlers/__init__.py b/unblob/handlers/__init__.py
@@ -1,6 +1,6 @@
 from ..models import Handlers
 from .archive import ar, arc, arj, cab, cpio, dmg, rar, sevenzip, stuffit, tar, zip
-from .compression import bzip2, compress, gzip, lz4, lzh, lzip, lzma, lzo, xz
+from .compression import bzip2, compress, gzip, lz4, lzh, lzip, lzma, lzo, xz, zstd
 from .executable import elf
 from .filesystem import (
     cramfs,
@@ -62,6 +62,7 @@
     lz4.SkippableFrameHandler,
     lz4.DefaultFrameHandler,
     xz.XZHandler,
+    zstd.ZSTDHandler,
     elf.ELF32Handler,
     elf.ELF64Handler,
 )
diff --git a/unblob/handlers/compression/zstd.py b/unblob/handlers/compression/zstd.py
@@ -0,0 +1,84 @@
+import io
+from typing import Optional
+
+from structlog import get_logger
+
+from unblob.extractors import Command
+
+from ...file_utils import Endian, InvalidInputFormat, convert_int8
+from ...models import File, Handler, HexString, ValidChunk
+
+logger = get_logger()
+
+MAGIC_LEN = 4
+BLOCK_HEADER_LEN = 3
+RAW_BLOCK = 0
+RLE_BLOCK = 1
+COMPRESSED_BLOCK = 2
+DICT_ID_FIELDSIZE_MAP = [0, 1, 2, 4]
+FRAME_CONTENT_FIELDSIZE_MAP = [0, 2, 4, 8]
+
+
+class ZSTDHandler(Handler):
+    NAME = "zstd"
+
+    PATTERNS = [HexString("28 B5 2F FD")]
+
+    EXTRACTOR = Command("zstd", "-d", "{inpath}", "-o", "{outdir}/{infile}")
+
+    def get_frame_header_size(self, frame_header_descriptor: int) -> int:
+        single_segment = (frame_header_descriptor >> 5 & 1) & 0b1
+        dictionary_id = frame_header_descriptor >> 0 & 0b11
+        frame_content_size = (frame_header_descriptor >> 6) & 0b1
+        return (
+            int(not single_segment)
+            + DICT_ID_FIELDSIZE_MAP[dictionary_id]
+            + FRAME_CONTENT_FIELDSIZE_MAP[frame_content_size]
+            + (single_segment and not frame_content_size)
+        )
+
+    def calculate_chunk(self, file: File, start_offset: int) -> Optional[ValidChunk]:
+
+        file.seek(start_offset, io.SEEK_SET)
+        file.seek(MAGIC_LEN, io.SEEK_CUR)
+
+        frame_header_descriptor = convert_int8(file.read(1), Endian.LITTLE)
+        frame_header_size = self.get_frame_header_size(frame_header_descriptor)
+
+        content_checksum_flag = frame_header_descriptor >> 2 & 1
+        if content_checksum_flag:
+            content_checksum_size = 4
+        else:
+            content_checksum_size = 0
+
+        unused_bit = frame_header_descriptor >> 4 & 1
+        reserved_bit = frame_header_descriptor >> 3 & 1
+
+        # these values MUST be zero per the standard
+        if unused_bit != 0x00 or reserved_bit != 0x0:
+            raise InvalidInputFormat("Invalid frame header format.")
+
+        file.seek(frame_header_size, io.SEEK_CUR)
+
+        last_block = False
+        while not last_block:
+            block_header = int.from_bytes(
+                file.read(BLOCK_HEADER_LEN), byteorder="little"
+            )
+            last_block = block_header >> 0 & 0b1
+            block_type = block_header >> 1 & 0b11
+
+            if block_type in [RAW_BLOCK, COMPRESSED_BLOCK]:
+                block_size = block_header >> 3
+            elif block_type == RLE_BLOCK:
+                block_size = 1
+            else:
+                raise InvalidInputFormat("Invalid block type")
+            file.seek(block_size, io.SEEK_CUR)
+
+        file.seek(content_checksum_size, io.SEEK_CUR)
+
+        return ValidChunk(
+            start_offset=start_offset,
+            end_offset=file.tell(),
+        )

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+version https://git-lfs.github.com/spec/v1`
	`2`	`+oid sha256:81468a5c3508970bed02d0824e394497407d0964bc37556de141347f1e320317`
	`3`	`+size 2632`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+version https://git-lfs.github.com/spec/v1`
	`2`	`+oid sha256:c79b72ddd77445e2ba70810ae04eef63e8dfb0dce9638200e418ea4a670b9f5c`
	`3`	`+size 2628`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+version https://git-lfs.github.com/spec/v1`
	`2`	`+oid sha256:37aab92d3ca37b75e87a30e7656d26ecd580b88a270cd09df8bcf7a0edd11177`
	`3`	`+size 2628`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+version https://git-lfs.github.com/spec/v1`
	`2`	`+oid sha256:c2619fb8093ec1eaab0669e4622a5ffe7bce9c28d03b74858a42b9c73c02b0b8`
	`3`	`+size 1312`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+version https://git-lfs.github.com/spec/v1`
	`2`	`+oid sha256:52738e82d35ccdb2c85e5dfcb30f8c320a7ae7e1b427203cd9e566f3d8b97820`
	`3`	`+size 2624`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+version https://git-lfs.github.com/spec/v1`
	`2`	`+oid sha256:36a4aba0dcec23e926f0692065524808479c33f4501bae79348f8f005f10b21c`
	`3`	`+size 4`