Add a zlib stream implementation

dissect/util/stream.py

@@ -1,5 +1,7 @@
 import io
 import os
+import sys
+import zlib
 from bisect import bisect_left, bisect_right
 from threading import Lock
 from typing import BinaryIO, Optional, Union
@@ -550,3 +552,93 @@ def _read(self, offset: int, length: int) -> bytes:
             overlay_idx += 1
 
         return b"".join(result)
+
+
+class ZlibStream(AlignedStream):
+    """Create a zlib stream from another file-like object.
+
+    Basically the same as ``gzip.GzipFile`` but for raw zlib streams.
+    Due to the nature of zlib streams, seeking backwards requires resetting the decompression context.
+
+    Args:
+        fh: The source file-like object.
+        size: The size the stream should be.
+    """
+
+    def __init__(self, fh: BinaryIO, size: Optional[int] = None, **kwargs):
+        self._fh = fh
+
+        self._zlib = None
+        self._zlib_args = kwargs
+        self._zlib_offset = 0
+        self._zlib_prepend = b""
+        self._zlib_prepend_offset = None
+        self._rewind()
+
+        super().__init__(size)
+
+    def _rewind(self) -> None:
+        self._fh.seek(0)
+        self._zlib = zlib.decompressobj(**self._zlib_args)
+        self._zlib_offset = 0
+        self._zlib_prepend = b""
+        self._zlib_prepend_offset = None
+
+    def _seek_zlib(self, offset: int) -> None:
+        if offset < self._zlib_offset:
+            self._rewind()
+
+        while self._zlib_offset < offset:
+            read_size = min(offset - self._zlib_offset, self.align)
+            if self._read_zlib(read_size) == b"":
+                break
+
+    def _read_fh(self, length: int) -> bytes:
+        if self._zlib_prepend_offset is None:
+            return self._fh.read(length)
+
+        if self._zlib_prepend_offset + length <= len(self._zlib_prepend):
+            offset = self._zlib_prepend_offset
+            self._zlib_prepend_offset += length
+            return self._zlib_prepend_offset[offset : self._zlib_prepend_offset]
+        else:
+            offset = self._zlib_prepend_offset
+            self._zlib_prepend_offset = None
+            return self._zlib_prepend[offset:] + self._fh.read(length - len(self._zlib_prepend) + offset)
+
+    def _read_zlib(self, length: int) -> bytes:
+        if length < 0:
+            return self.readall()
+
+        result = []
+        while length > 0:
+            buf = self._read_fh(io.DEFAULT_BUFFER_SIZE)
+            decompressed = self._zlib.decompress(buf, length)
+
+            if self._zlib.unconsumed_tail != b"":
+                self._zlib_prepend = self._zlib.unconsumed_tail
+                self._zlib_prepend_offset = 0
+
+            if buf == b"":
+                break
+
+            result.append(decompressed)
+            length -= len(decompressed)
+
+        buf = b"".join(result)
+        self._zlib_offset += len(buf)
+        return buf
+
+    def _read(self, offset: int, length: int) -> bytes:
+        self._seek_zlib(offset)
+        return self._read_zlib(length)
+
+    def readall(self) -> bytes:
+        chunks = []
+        # sys.maxsize means the max length of output buffer is unlimited,
+        # so that the whole input buffer can be decompressed within one
+        # .decompress() call.
+        while data := self._read_zlib(sys.maxsize):
+            chunks.append(data)
+
+        return b"".join(chunks)

tests/test_stream.py

@@ -1,4 +1,5 @@
 import io
+import zlib
 
 import pytest
 
@@ -170,3 +171,23 @@ def test_overlay_stream():
     fh.add((512 * 8) - 4, b"\x04" * 100)
     fh.seek((512 * 8) - 4)
     assert fh.read(100) == b"\x04" * 4
+
+
+def test_zlib_stream():
+    buf = io.BytesIO(zlib.compress(b"\x01" * 512 + b"\x02" * 512 + b"\x03" * 512 + b"\x04" * 512))
+    fh = stream.ZlibStream(buf, size=512 * 4)
+
+    assert fh.read(512) == b"\x01" * 512
+    assert fh.read(512) == b"\x02" * 512
+    assert fh.read(512) == b"\x03" * 512
+    assert fh.read(512) == b"\x04" * 512
+    assert fh.read(1) == b""
+
+    fh.seek(0)
+    assert fh.read(512) == b"\x01" * 512
+
+    fh.seek(1024)
+    assert fh.read(512) == b"\x03" * 512
+
+    fh.seek(512)
+    assert fh.read(1024) == b"\x02" * 512 + b"\x03" * 512