From 3256f431f0bb09afa170787c9e6ebf784f811ee7 Mon Sep 17 00:00:00 2001 From: joel Date: Thu, 13 Oct 2022 15:50:01 -0400 Subject: [PATCH 1/4] add zstd packer/unpacker and test --- ofrak_components/Dockerstub | 3 +- ofrak_components/ofrak_components/zstd.py | 91 +++++++++++++++++++ .../test_zstd_component.py | 45 +++++++++ 3 files changed, 138 insertions(+), 1 deletion(-) create mode 100644 ofrak_components/ofrak_components/zstd.py create mode 100644 ofrak_components/ofrak_components_test/test_zstd_component.py diff --git a/ofrak_components/Dockerstub b/ofrak_components/Dockerstub index 97330fcdd..c11c8a1a6 100644 --- a/ofrak_components/Dockerstub +++ b/ofrak_components/Dockerstub @@ -14,7 +14,8 @@ RUN apt-get -y update && \ qemu \ qemu-user-static \ u-boot-tools \ - unar + unar \ + zstd # Install apktool and uber-apk-signer RUN apt-get -y update && apt-get -y install openjdk-11-jdk diff --git a/ofrak_components/ofrak_components/zstd.py b/ofrak_components/ofrak_components/zstd.py new file mode 100644 index 000000000..b28c1e66d --- /dev/null +++ b/ofrak_components/ofrak_components/zstd.py @@ -0,0 +1,91 @@ +import subprocess +import tempfile +from dataclasses import dataclass + +from ofrak import Packer, Unpacker, Resource +from ofrak.component.packer import PackerError +from ofrak.component.unpacker import UnpackerError +from ofrak.core import ( + GenericBinary, + format_called_process_error, + MagicMimeIdentifier, + MagicDescriptionIdentifier, +) +from ofrak.model.component_model import CC, ComponentConfig +from ofrak_type.range import Range + + +class ZstdData(GenericBinary): + """ + A zstd binary blob. + """ + + async def get_child(self) -> GenericBinary: + return await self.resource.get_only_child_as_view(GenericBinary) + +@dataclass +class ZstdPackerConfig(ComponentConfig): + compression_level: int + + +class ZstdUnpacker(Unpacker[None]): + """ + Unpack (decompress) a zstd file. + """ + + id = b"ZstdUnpacker" + targets = (ZstdData,) + children = (GenericBinary,) + + async def unpack(self, resource: Resource, config: CC) -> None: + with tempfile.NamedTemporaryFile(suffix=".zstd") as compressed_file: + compressed_file.write(await resource.get_data()) + compressed_file.flush() + output_filename = tempfile.mktemp() + + command = ["zstd", "-d", "-k", compressed_file.name, "-o", output_filename] + try: + subprocess.run(command, check=True) + with open(output_filename, 'rb') as f: + result = f.read() + except subprocess.CalledProcessError as e: + raise UnpackerError(format_called_process_error(e)) + + await resource.create_child(tags=(GenericBinary,), data=result) + + +class ZstdPacker(Packer[None]): + """ + Pack data into a compressed zstd file. + """ + + targets = (ZstdData,) + + async def pack(self, resource: Resource, config: ZstdPackerConfig = ZstdPackerConfig(compression_level=19)): + zstd_view = await resource.view_as(ZstdData) + child_file = await zstd_view.get_child() + uncompressed_data = await child_file.resource.get_data() + + with tempfile.NamedTemporaryFile() as uncompressed_file: + uncompressed_file.write(uncompressed_data) + uncompressed_file.flush() + output_filename = tempfile.mktemp() + + command = ["zstd", "-T0", f"-{config.compression_level}"] + if config.compression_level > 19: + command.append("--ultra") + command.extend([uncompressed_file.name, "-o", output_filename]) + try: + subprocess.run(command, check=True) + with open(output_filename, 'rb') as f: + result = f.read() + except subprocess.CalledProcessError as e: + raise PackerError(format_called_process_error(e)) + + compressed_data = result + original_size = await zstd_view.resource.get_data_length() + resource.queue_patch(Range(0, original_size), compressed_data) + + +MagicMimeIdentifier.register(ZstdData, "application/x-zstd") +MagicDescriptionIdentifier.register(ZstdData, lambda s: s.lower().startswith("zstandard compressed data")) diff --git a/ofrak_components/ofrak_components_test/test_zstd_component.py b/ofrak_components/ofrak_components_test/test_zstd_component.py new file mode 100644 index 000000000..16759f9cb --- /dev/null +++ b/ofrak_components/ofrak_components_test/test_zstd_component.py @@ -0,0 +1,45 @@ +import subprocess +import tempfile + +import pytest + +from ofrak.core.filesystem import format_called_process_error +from ofrak.resource import Resource +from pytest_ofrak.patterns.compressed_filesystem_unpack_modify_pack import ( + CompressedFileUnpackModifyPackPattern, +) + + +class TestZstdUnpackModifyPack(CompressedFileUnpackModifyPackPattern): + @pytest.fixture(autouse=True) + def create_test_file(self, tmpdir): + d = tmpdir.mkdir("zstd") + uncompressed_filename = d.join("hello.txt").realpath() + with open(uncompressed_filename, "wb") as f: + f.write(self.INITIAL_DATA) + + compressed_filename = d.join("hello.zstd").realpath() + command = ["zstd", "-19", uncompressed_filename, "-o", compressed_filename] + try: + subprocess.run(command, check=True, capture_output=True) + except subprocess.CalledProcessError as e: + raise RuntimeError(format_called_process_error(e)) + + self._test_file = compressed_filename + + async def verify(self, repacked_root_resource: Resource) -> None: + compressed_data = await repacked_root_resource.get_data() + with tempfile.NamedTemporaryFile(suffix=".zstd") as compressed_file: + compressed_file.write(compressed_data) + compressed_file.flush() + output_filename = tempfile.mktemp() + + command = ["zstd", "-d", "-k", compressed_file.name, "-o", output_filename] + try: + subprocess.run(command, check=True, capture_output=True) + with open(output_filename, 'rb') as f: + result = f.read() + except subprocess.CalledProcessError as e: + raise RuntimeError(format_called_process_error(e)) + + assert result == self.EXPECTED_REPACKED_DATA From f571cf9b6d15bd980d6f77a85fc292a67237bcde Mon Sep 17 00:00:00 2001 From: joel Date: Thu, 13 Oct 2022 16:57:16 -0400 Subject: [PATCH 2/4] run black --- ofrak_components/ofrak_components/zstd.py | 13 +++++++++---- .../ofrak_components_test/test_zstd_component.py | 2 +- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/ofrak_components/ofrak_components/zstd.py b/ofrak_components/ofrak_components/zstd.py index b28c1e66d..27b9ee7be 100644 --- a/ofrak_components/ofrak_components/zstd.py +++ b/ofrak_components/ofrak_components/zstd.py @@ -23,6 +23,7 @@ class ZstdData(GenericBinary): async def get_child(self) -> GenericBinary: return await self.resource.get_only_child_as_view(GenericBinary) + @dataclass class ZstdPackerConfig(ComponentConfig): compression_level: int @@ -46,7 +47,7 @@ async def unpack(self, resource: Resource, config: CC) -> None: command = ["zstd", "-d", "-k", compressed_file.name, "-o", output_filename] try: subprocess.run(command, check=True) - with open(output_filename, 'rb') as f: + with open(output_filename, "rb") as f: result = f.read() except subprocess.CalledProcessError as e: raise UnpackerError(format_called_process_error(e)) @@ -61,7 +62,9 @@ class ZstdPacker(Packer[None]): targets = (ZstdData,) - async def pack(self, resource: Resource, config: ZstdPackerConfig = ZstdPackerConfig(compression_level=19)): + async def pack( + self, resource: Resource, config: ZstdPackerConfig = ZstdPackerConfig(compression_level=19) + ): zstd_view = await resource.view_as(ZstdData) child_file = await zstd_view.get_child() uncompressed_data = await child_file.resource.get_data() @@ -77,7 +80,7 @@ async def pack(self, resource: Resource, config: ZstdPackerConfig = ZstdPackerCo command.extend([uncompressed_file.name, "-o", output_filename]) try: subprocess.run(command, check=True) - with open(output_filename, 'rb') as f: + with open(output_filename, "rb") as f: result = f.read() except subprocess.CalledProcessError as e: raise PackerError(format_called_process_error(e)) @@ -88,4 +91,6 @@ async def pack(self, resource: Resource, config: ZstdPackerConfig = ZstdPackerCo MagicMimeIdentifier.register(ZstdData, "application/x-zstd") -MagicDescriptionIdentifier.register(ZstdData, lambda s: s.lower().startswith("zstandard compressed data")) +MagicDescriptionIdentifier.register( + ZstdData, lambda s: s.lower().startswith("zstandard compressed data") +) diff --git a/ofrak_components/ofrak_components_test/test_zstd_component.py b/ofrak_components/ofrak_components_test/test_zstd_component.py index 16759f9cb..d5035ddf5 100644 --- a/ofrak_components/ofrak_components_test/test_zstd_component.py +++ b/ofrak_components/ofrak_components_test/test_zstd_component.py @@ -37,7 +37,7 @@ async def verify(self, repacked_root_resource: Resource) -> None: command = ["zstd", "-d", "-k", compressed_file.name, "-o", output_filename] try: subprocess.run(command, check=True, capture_output=True) - with open(output_filename, 'rb') as f: + with open(output_filename, "rb") as f: result = f.read() except subprocess.CalledProcessError as e: raise RuntimeError(format_called_process_error(e)) From ec386b2a8b654b6722d35aac832df5c8b9d25834 Mon Sep 17 00:00:00 2001 From: Jacob Strieb <99368685+rbs-jacob@users.noreply.github.com> Date: Tue, 18 Oct 2022 20:06:10 -0400 Subject: [PATCH 3/4] Apply suggestions from code review --- ofrak_components/ofrak_components/zstd.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ofrak_components/ofrak_components/zstd.py b/ofrak_components/ofrak_components/zstd.py index 27b9ee7be..7cd8e6f52 100644 --- a/ofrak_components/ofrak_components/zstd.py +++ b/ofrak_components/ofrak_components/zstd.py @@ -55,16 +55,16 @@ async def unpack(self, resource: Resource, config: CC) -> None: await resource.create_child(tags=(GenericBinary,), data=result) -class ZstdPacker(Packer[None]): +class ZstdPacker(Packer[ZstdPackerConfig]): """ Pack data into a compressed zstd file. """ targets = (ZstdData,) - async def pack( - self, resource: Resource, config: ZstdPackerConfig = ZstdPackerConfig(compression_level=19) - ): + async def pack(self, resource: Resource, config: Optional[ZstdPackerConfig] = None): + if config is None: + config = ZstdPackerConfig(compression_level=19) zstd_view = await resource.view_as(ZstdData) child_file = await zstd_view.get_child() uncompressed_data = await child_file.resource.get_data() From 157024d490391861ada280497c8d6e0e6dc16a78 Mon Sep 17 00:00:00 2001 From: Jacob Strieb <99368685+rbs-jacob@users.noreply.github.com> Date: Tue, 18 Oct 2022 20:45:21 -0400 Subject: [PATCH 4/4] Update ofrak_components/ofrak_components/zstd.py --- ofrak_components/ofrak_components/zstd.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ofrak_components/ofrak_components/zstd.py b/ofrak_components/ofrak_components/zstd.py index 7cd8e6f52..386e38e2a 100644 --- a/ofrak_components/ofrak_components/zstd.py +++ b/ofrak_components/ofrak_components/zstd.py @@ -1,6 +1,7 @@ import subprocess import tempfile from dataclasses import dataclass +from typing import Optional from ofrak import Packer, Unpacker, Resource from ofrak.component.packer import PackerError