diff --git a/dissect/hypervisor/disk/qcow2.py b/dissect/hypervisor/disk/qcow2.py index e869f71..a81ecfb 100644 --- a/dissect/hypervisor/disk/qcow2.py +++ b/dissect/hypervisor/disk/qcow2.py @@ -3,9 +3,9 @@ # - https://github.com/qemu/qemu/blob/master/docs/interop/qcow2.txt from __future__ import annotations +import sys import zlib from functools import cached_property, lru_cache -from io import BytesIO from pathlib import Path from typing import TYPE_CHECKING, BinaryIO @@ -28,8 +28,10 @@ from collections.abc import Iterator try: - import zstandard as zstd - + if sys.version_info >= (3, 14): + from compression import zstd # novermin + else: + from backports import zstd HAS_ZSTD = True except ImportError: HAS_ZSTD = False @@ -384,16 +386,8 @@ def _decompress(self, buf: bytes) -> bytes: return dctx.decompress(buf, self.qcow2.cluster_size) if self.qcow2.compression_type == c_qcow2.QCOW2_COMPRESSION_TYPE_ZSTD: - result = [] - dctx = zstd.ZstdDecompressor() - reader = dctx.stream_reader(BytesIO(buf)) - while reader.tell() < self.qcow2.cluster_size: - chunk = reader.read(self.qcow2.cluster_size - reader.tell()) - if not chunk: - break - result.append(chunk) - return b"".join(result) + return dctx.decompress(buf, self.qcow2.cluster_size) raise Error(f"Invalid compression type: {self.qcow2.compression_type}") diff --git a/pyproject.toml b/pyproject.toml index 8404869..3ea276d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ repository = "https://github.com/fox-it/dissect.hypervisor" [project.optional-dependencies] full = [ "pycryptodome", + "backports.zstd; python_version < '3.14'", ] dev = [ "dissect.hypervisor[full]", diff --git a/tests/_data/disk/qcow2/basic-zstd.qcow2.gz b/tests/_data/disk/qcow2/basic-zstd.qcow2.gz new file mode 100644 index 0000000..33d3cca Binary files /dev/null and b/tests/_data/disk/qcow2/basic-zstd.qcow2.gz differ diff --git a/tests/conftest.py b/tests/conftest.py index 9bb7818..3673181 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -89,6 +89,11 @@ def basic_qcow2() -> Iterator[BinaryIO]: yield from open_file_gz("_data/disk/qcow2/basic.qcow2.gz") +@pytest.fixture +def basic_zstd_qcow2() -> Iterator[BinaryIO]: + yield from open_file_gz("_data/disk/qcow2/basic-zstd.qcow2.gz") + + @pytest.fixture def data_file_qcow2() -> Path: return absolute_path("_data/disk/qcow2/data-file.qcow2.gz") diff --git a/tests/disk/test_qcow2.py b/tests/disk/test_qcow2.py index e7d6e42..017a919 100644 --- a/tests/disk/test_qcow2.py +++ b/tests/disk/test_qcow2.py @@ -16,8 +16,9 @@ def mock_open_gz(self: Path, *args, **kwargs) -> BinaryIO: return gzip.open(self if self.suffix.lower() == ".gz" else self.with_suffix(self.suffix + ".gz")) -def test_basic(basic_qcow2: BinaryIO) -> None: - qcow2 = QCow2(basic_qcow2) +@pytest.mark.parametrize("name", ["basic_qcow2", "basic_zstd_qcow2"]) +def test_basic(name: str, request: pytest.FixtureRequest) -> None: + qcow2 = QCow2(request.getfixturevalue(name)) assert qcow2.backing_file is None assert qcow2.data_file is qcow2.fh