Reset file position in container, volume and filesystem (fox-it#437)

Author: Schamper

dissect/target/container.py

@@ -65,13 +65,40 @@ def detect(cls, item: Union[list, BinaryIO, Path]) -> bool:
 
         return False
 
-    @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+    @classmethod
+    def detect_fh(cls, fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
         """Detect if this ``Container`` can be used to open the file-like object ``fh``.
 
-        The function checks wether the raw data contains any magic information that corresponds to this
+        The function checks whether the raw data contains any magic information that corresponds to this
         specific container.
 
+        Args:
+            fh: A file-like object that we want to open a ``Container`` on.
+            original: The original argument passed to ``detect()``.
+
+        Returns:
+            ``True`` if this ``Container`` can be used for this file-like object, ``False`` otherwise.
+        """
+        offset = fh.tell()
+        try:
+            fh.seek(0)
+            return cls._detect_fh(fh, original)
+        except NotImplementedError:
+            raise
+        except Exception as e:
+            log.warning("Failed to detect %s container", cls.__name__)
+            log.debug("", exc_info=e)
+        finally:
+            fh.seek(offset)
+
+        return False
+
+    @staticmethod
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+        """Detect if this ``Container`` can be used to open the file-like object ``fh``.
+
+        This method should be implemented by subclasses. The position of ``fh`` is guaranteed to be ``0``.
+
         Args:
             fh: A file-like object that we want to open a ``Container`` on.
             original: The original argument passed to ``detect()``.
@@ -177,6 +204,7 @@ def open(item: Union[list, str, BinaryIO, Path], *args, **kwargs):
     first = item[0] if isinstance(item, list) else item
     first_fh = None
     first_fh_opened = False
+    first_fh_offset = None
     first_path = None
 
     if hasattr(first, "read"):
@@ -187,6 +215,10 @@ def open(item: Union[list, str, BinaryIO, Path], *args, **kwargs):
             first_fh = first.open("rb")
             first_fh_opened = True
 
+    if first_fh:
+        first_fh_offset = first_fh.tell()
+        first_fh.seek(0)
+
     try:
         for container in CONTAINERS + [RawContainer]:
             try:
@@ -203,6 +235,8 @@ def open(item: Union[list, str, BinaryIO, Path], *args, **kwargs):
     finally:
         if first_fh_opened:
             first_fh.close()
+        elif first_fh:
+            first_fh.seek(first_fh_offset)
 
     raise ContainerError(f"Failed to detect container for {item}")
 

dissect/target/containers/asdf.py

@@ -26,11 +26,8 @@ def __init__(self, fh: BinaryIO, *args, **kwargs):
         super().__init__(fh, self.asdf.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
-        magic = fh.read(4)
-        fh.seek(-4, io.SEEK_CUR)
-
-        return magic == FILE_MAGIC
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+        return fh.read(4) == FILE_MAGIC
 
     @staticmethod
     def detect_path(path: Path, original: Union[list, BinaryIO]) -> bool:

dissect/target/containers/ewf.py

@@ -23,12 +23,9 @@ def __init__(self, fh: Union[list, BinaryIO, Path], *args, **kwargs):
         super().__init__(fh, self.ewf.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
         """Detect file header"""
-        magic = fh.read(3)
-        fh.seek(-3, io.SEEK_CUR)
-
-        return magic == b"EVF" or magic == b"LVF" or magic == b"LEF"
+        return fh.read(3) in (b"EVF", b"LVF", b"LEF")
 
     @staticmethod
     def detect_path(path: Path, original: Union[list, BinaryIO]) -> bool:

dissect/target/containers/hdd.py

@@ -17,7 +17,7 @@ def __init__(self, fh: Path, *args, **kwargs):
         super().__init__(fh, self.stream.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
         return False
 
     @staticmethod

dissect/target/containers/hds.py

@@ -18,11 +18,8 @@ def __init__(self, fh: Union[BinaryIO, Path], *args, **kwargs):
         super().__init__(fh, self.hds.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
-        sig = fh.read(16)
-        fh.seek(-16, io.SEEK_CUR)
-
-        return sig in (c_hdd.SIGNATURE_STRUCTURED_DISK_V1, c_hdd.SIGNATURE_STRUCTURED_DISK_V2)
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+        return fh.read(16) in (c_hdd.SIGNATURE_STRUCTURED_DISK_V1, c_hdd.SIGNATURE_STRUCTURED_DISK_V2)
 
     @staticmethod
     def detect_path(path: Path, original: Union[list, BinaryIO]) -> bool:

dissect/target/containers/qcow2.py

@@ -17,11 +17,8 @@ def __init__(self, fh: Union[BinaryIO, Path], data_file=None, backing_file=None,
         super().__init__(fh, self.qcow2.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
-        magic = fh.read(4)
-        fh.seek(-4, io.SEEK_CUR)
-
-        return magic == c_qcow2.QCOW2_MAGIC_BYTES
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+        return fh.read(4) == c_qcow2.QCOW2_MAGIC_BYTES
 
     @staticmethod
     def detect_path(path: Path, original: Union[list, BinaryIO]) -> bool:

dissect/target/containers/raw.py

@@ -32,7 +32,7 @@ def __init__(self, fh: Union[BinaryIO, Path], *args, **kwargs):
         super().__init__(fh, size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
         return True
 
     @staticmethod

dissect/target/containers/split.py

@@ -32,7 +32,7 @@ def __init__(self, fh: Union[list, BinaryIO, Path], *args, **kwargs):
         super().__init__(fh, offset, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
         return isinstance(original, list) and len(original) > 1
 
     @staticmethod

dissect/target/containers/vdi.py

@@ -19,10 +19,8 @@ def __init__(self, fh: Union[BinaryIO, Path], *args, **kwargs):
         super().__init__(fh, self.vdi.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
-        magic = fh.read(68)
-        fh.seek(-68, io.SEEK_CUR)
-        return magic[-4:] == b"\x7f\x10\xda\xbe"
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+        return fh.read(68)[-4:] == b"\x7f\x10\xda\xbe"
 
     @staticmethod
     def detect_path(path: Path, original: Union[list, BinaryIO]) -> bool:

dissect/target/containers/vhd.py

@@ -17,13 +17,9 @@ def __init__(self, fh: Union[BinaryIO, Path], *args, **kwargs):
         super().__init__(fh, self.vhd.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
-        offset = fh.tell()
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
         fh.seek(-512, io.SEEK_END)
-        magic = fh.read(9)
-        fh.seek(offset)
-
-        return b"conectix" in magic
+        return b"conectix" in fh.read(9)
 
     @staticmethod
     def detect_path(path: Path, original: Union[list, BinaryIO]) -> bool:

dissect/target/containers/vhdx.py

@@ -13,11 +13,8 @@ def __init__(self, fh: Union[BinaryIO, Path], *args, **kwargs):
         super().__init__(fh, self.vhdx.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
-        magic = fh.read(8)
-        fh.seek(-8, io.SEEK_CUR)
-
-        return magic == b"vhdxfile"
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+        return fh.read(8) == b"vhdxfile"
 
     @staticmethod
     def detect_path(path: Path, original: Union[list, BinaryIO]) -> bool:

dissect/target/containers/vmdk.py

@@ -15,11 +15,8 @@ def __init__(self, fh: Union[BinaryIO, Path], *args, **kwargs):
         super().__init__(fh, self.vmdk.size, *args, **kwargs)
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
-        magic = fh.read(4)
-        fh.seek(-4, io.SEEK_CUR)
-
-        return magic in (b"KDMV", b"COWD", b"# Di")
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+        return fh.read(4) in (b"KDMV", b"COWD", b"# Di")
 
     @staticmethod
     def detect_path(path: Path, original: Union[list, BinaryIO]) -> bool:

dissect/target/filesystem.py

@@ -1449,15 +1449,22 @@ def register(module: str, class_name: str, internal: bool = True) -> None:
 
 
 def open(fh: BinaryIO, *args, **kwargs) -> Filesystem:
-    for filesystem in FILESYSTEMS:
-        try:
-            if filesystem.detect(fh):
-                instance = filesystem(fh, *args, **kwargs)
-                instance.volume = fh
-                return instance
-        except ImportError as e:
-            log.info("Failed to import %s", filesystem)
-            log.debug("", exc_info=e)
+    offset = fh.tell()
+    fh.seek(0)
+
+    try:
+        for filesystem in FILESYSTEMS:
+            try:
+                if filesystem.detect(fh):
+                    instance = filesystem(fh, *args, **kwargs)
+                    instance.volume = fh
+
+                    return instance
+            except ImportError as e:
+                log.info("Failed to import %s", filesystem)
+                log.debug("", exc_info=e)
+    finally:
+        fh.seek(offset)
 
     raise FilesystemError(f"Failed to open filesystem for {fh}")
 

dissect/target/volume.py

@@ -327,10 +327,15 @@ def open(fh: BinaryIO, *args, **kwargs) -> DissectVolumeSystem:
     Returns:
         An opened :class:`~dissect.target.volumes.disk.DissectVolumeSystem`.
     """
+    offset = fh.tell()
+    fh.seek(0)
+
     try:
         return disk.DissectVolumeSystem(fh)
     except Exception as e:
         raise VolumeSystemError(f"Failed to load volume system for {fh}", cause=e)
+    finally:
+        fh.seek(offset)
 
 
 def is_lvm_volume(volume: BinaryIO) -> bool:

tests/data/plugin_register/container.py

@@ -13,7 +13,7 @@ def __repr__(self):
         return f"<{self.__class__.__name__} size={self.size} vs={self.vs}>"
 
     @staticmethod
-    def detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
+    def _detect_fh(fh: BinaryIO, original: Union[list, BinaryIO]) -> bool:
         return False
 
     @staticmethod

tests/test_container_open.py

@@ -1,4 +1,5 @@
 import struct
+from io import BytesIO
 from pathlib import Path
 from unittest.mock import Mock, patch
 
@@ -69,3 +70,28 @@ def test_open_fallback_fh(tmp_path):
     tmp_dummy.write_bytes(b"\x00" * 1024)
     assert isinstance(container.open(tmp_dummy), raw.RawContainer)
     assert not vhd.VhdContainer.detect(tmp_dummy)
+
+
+def test_reset_file_position() -> None:
+    fh = BytesIO(b"\x00" * 8192)
+    fh.seek(512)
+
+    class MockContainer(container.Container):
+        def __init__(self, fh):
+            assert fh.tell() == 0
+            fh.seek(1024)
+            self.success = True
+
+        @staticmethod
+        def _detect_fh(fh, *args, **kwargs):
+            assert fh.tell() == 0
+            fh.seek(256)
+            return True
+
+    mock_container = Mock()
+    mock_container.MockContainer = MockContainer
+    with patch.object(container, "CONTAINERS", [mock_container.MockContainer]):
+        opened_container = container.open(fh)
+        assert isinstance(opened_container, mock_container.MockContainer)
+        assert opened_container.success
+        assert fh.tell() == 512

tests/test_filesystem.py

@@ -1,12 +1,14 @@
 import os
 import stat
+from io import BytesIO
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 from typing import Union
 from unittest.mock import Mock, patch
 
 import pytest
 from _pytest.fixtures import FixtureRequest
 
+from dissect.target import filesystem
 from dissect.target.exceptions import (
     FileNotFoundError,
     NotADirectoryError,
@@ -1092,3 +1094,29 @@ def test_mapped_file_lattr(mapped_file: MappedFile) -> None:
     with patch("dissect.target.helpers.fsutil.fs_attrs", autospec=True) as fs_attrs:
         mapped_file.lattr()
         fs_attrs.assert_called_with(mapped_file.entry, follow_symlinks=False)
+
+
+def test_reset_file_position() -> None:
+    fh = BytesIO(b"\x00" * 8192)
+    fh.seek(512)
+
+    class MockFilesystem(filesystem.Filesystem):
+        def __init__(self, fh):
+            assert fh.tell() == 0
+            fh.seek(1024)
+            self.success = True
+
+        @staticmethod
+        def _detect(fh):
+            assert fh.tell() == 0
+            fh.seek(256)
+            return True
+
+    mock_fs = Mock()
+    mock_fs.MockFilesystem = MockFilesystem
+
+    with patch.object(filesystem, "FILESYSTEMS", [mock_fs.MockFilesystem]):
+        opened_fs = filesystem.open(fh)
+        assert isinstance(opened_fs, mock_fs.MockFilesystem)
+        assert opened_fs.success
+        assert fh.tell() == 512

tests/test_volume.py

@@ -0,0 +1,31 @@
+from io import BytesIO
+from unittest.mock import patch
+
+from dissect.target import volume
+from dissect.target.volumes import disk
+
+
+def test_reset_file_position() -> None:
+    fh = BytesIO(b"\x00" * 8192)
+    fh.seek(512)
+
+    class MockVolumeSystem(volume.VolumeSystem):
+        def __init__(self, fh):
+            assert fh.tell() == 0
+            fh.seek(1024)
+            self.success = True
+
+        @staticmethod
+        def _detect(fh):
+            assert fh.tell() == 0
+            fh.seek(256)
+            return True
+
+    with patch.object(disk, "DissectVolumeSystem", MockVolumeSystem):
+        assert MockVolumeSystem.detect(fh)
+        assert fh.tell() == 512
+
+        opened_vs = volume.open(fh)
+        assert isinstance(opened_vs, MockVolumeSystem)
+        assert opened_vs.success
+        assert fh.tell() == 512