Address review comments

Author: Schamper

dissect/hypervisor/backup/vbk.py

@@ -5,22 +5,17 @@
 
 from functools import cached_property, lru_cache
 from io import BytesIO
-from typing import BinaryIO, Generic, Iterator, Optional, TypeVar
+from typing import BinaryIO, Generic, Iterator, TypeVar
 from zlib import crc32
 
 from dissect.cstruct import Structure
+from dissect.hypervisor.backup.c_vbk import c_vbk
+from dissect.hypervisor.exceptions import Error
+from dissect.util.compression import lz4
 from dissect.util.crc32c import crc32c
 from dissect.util.stream import AlignedStream
 from dissect.util.xmemoryview import xmemoryview
 
-try:
-    from lz4.block import decompress as lz4_decompress
-except ImportError:
-    from dissect.util.compression.lz4 import decompress as lz4_decompress
-
-from dissect.hypervisor.backup.c_vbk import c_vbk
-from dissect.hypervisor.exceptions import Error
-
 PAGE_SIZE = 4096
 """VBK page size."""
 
@@ -108,7 +103,7 @@ def read_meta_blob(self, page: int) -> bytes:
         """
         return self.active_slot._get_meta_blob(page)
 
-    def get(self, path: str, item: Optional[DirItem] = None) -> DirItem:
+    def get(self, path: str, item: DirItem | None = None) -> DirItem:
         """Get a directory item from the VBK file."""
         item = item or self.root
 
@@ -351,7 +346,7 @@ def size(self) -> int:
         raise VBKError(f"Size not available for {self!r}")
 
     @cached_property
-    def properties(self) -> Optional[PropertiesDictionary]:
+    def properties(self) -> PropertiesDictionary | None:
         """The properties of the directory item, if it has them."""
         if self.entry.PropsRootPage == -1:
             return None
@@ -360,25 +355,22 @@ def properties(self) -> Optional[PropertiesDictionary]:
 
     def is_dir(self) -> bool:
         """Return whether the directory item is a directory."""
-        return isinstance(self, (RootDirectory, SubFolderItem))
+        return False
 
     def is_file(self) -> bool:
         """Return whether the directory item is a file."""
         return self.is_internal_file() or self.is_external_file()
 
     def is_internal_file(self) -> bool:
         """Return whether the directory item is an internal file."""
-        return isinstance(self, IntFibItem)
+        return False
 
     def is_external_file(self) -> bool:
         """Return whether the directory item is an external file."""
-        return isinstance(self, ExtFibItem)
+        return False
 
     def listdir(self) -> dict[str, DirItem]:
         """Return a dictionary of the items in the directory."""
-        if not self.is_dir():
-            raise NotADirectoryError(f"{self!r} is not a directory")
-
         return {item.name: item for item in self.iterdir()}
 
     def iterdir(self) -> Iterator[DirItem]:
@@ -402,6 +394,9 @@ def __init__(self, vbk: VBK, page: int, count: int):
     def __repr__(self) -> str:
         return f"<RootDirectory root={self.root} count={self.count}>"
 
+    def is_dir(self) -> bool:
+        return True
+
     def iterdir(self) -> Iterator[DirItem]:
         yield from MetaVector(self.vbk, DirItem, self.root, self.count)
 
@@ -422,6 +417,9 @@ def __init__(self, vbk: VBK, buf: bytes):
     def __repr__(self) -> str:
         return f"<SubFolderItem name={self.name!r} root={self.root} count={self.count}>"
 
+    def is_dir(self) -> bool:
+        return True
+
     def iterdir(self) -> Iterator[DirItem]:
         yield from MetaVector(self.vbk, DirItem, self.root, self.count)
 
@@ -441,6 +439,9 @@ def __repr__(self) -> str:
     def size(self) -> int:
         return self.entry.ExtFib.FibSize
 
+    def is_external_file(self) -> bool:
+        return True
+
 
 class IntFibItem(DirItem):
     """Directory item for an internal file.
@@ -460,6 +461,9 @@ def __repr__(self) -> str:
     def size(self) -> int:
         return self.entry.IntFib.FibSize
 
+    def is_internal_file(self) -> bool:
+        return True
+
     def open(self) -> FibStream:
         return FibStream(
             self.vbk,
@@ -477,13 +481,13 @@ class PatchItem(DirItem):
         - **TODO**: CPatchMeta
     """
 
+    def __repr__(self) -> str:
+        return f"<PatchItem name={self.name!r} size={self.size}>"
+
     @cached_property
     def size(self) -> int:
         return self.entry.Patch.FibSize
 
-    def __repr__(self) -> str:
-        return f"<PatchItem name={self.name!r} size={self.size}>"
-
 
 class IncrementItem(DirItem):
     """Directory item for an increment.
@@ -493,13 +497,13 @@ class IncrementItem(DirItem):
         - **TODO**: CIncrementMeta
     """
 
+    def __repr__(self) -> str:
+        return f"<IncrementItem name={self.name!r} size={self.size}>"
+
     @cached_property
     def size(self) -> int:
         return self.entry.Increment.FibSize
 
-    def __repr__(self) -> str:
-        return f"<IncrementItem name={self.name!r} size={self.size}>"
-
 
 class MetaTableDescriptor(MetaItem):
     """A descriptor for a meta table in the VBK file.
@@ -1030,7 +1034,7 @@ def _read(self, offset: int, length: int) -> bytes:
                 if block.is_compressed():
                     if block.compression_type == c_vbk.CompressionType.LZ4:
                         # First 12 bytes are Lz4BlockHeader
-                        buf = lz4_decompress(memoryview(buf)[12:], block.source_size)
+                        buf = lz4.decompress(memoryview(buf)[12:], block.source_size)
                     else:
                         raise VBKError(f"Unsupported compression type: {block.compression_type}")
 

dissect/hypervisor/tools/backup.py

@@ -4,7 +4,7 @@
 from pathlib import Path
 
 from dissect.hypervisor.backup.c_vma import c_vma
-from dissect.hypervisor.backup.vbk import VBK
+from dissect.hypervisor.backup.vbk import VBK, DirItem
 from dissect.hypervisor.backup.vma import VMA, _iter_mask
 
 try:
@@ -144,20 +144,22 @@ def extract_vma(vma: VMA, out_dir: Path) -> None:
 
 
 def extract_vbk(vbk: VBK, out_dir: Path) -> None:
-    with progress:
-        try:
-            root_dir = next(vbk.get("/").iterdir())
-            out_dir = out_dir.joinpath(root_dir.name)
-            out_dir.mkdir(exist_ok=True)
-
-            for entry in root_dir.iterdir():
-                out_file = out_dir.joinpath(entry.name)
+    def extract_directory(directory: DirItem, out_dir: Path) -> None:
+        out_dir.mkdir(exist_ok=True)
+        for entry in directory.iterdir():
+            out_path = out_dir.joinpath(entry.name)
+            if entry.is_dir():
+                extract_directory(entry, out_path)
+            else:
                 task_id = progress.add_task("extract", filename=entry.name, total=entry.size)
-
-                with entry.open() as fh_in, out_file.open("wb") as fh_out:
+                with entry.open() as fh_in, out_path.open("wb") as fh_out:
                     for chunk in iter(lambda: fh_in.read(vbk.block_size), b""):
                         fh_out.write(chunk)
                         progress.update(task_id, advance=len(chunk))
+
+    with progress:
+        try:
+            extract_directory(vbk.get("/"), out_dir)
         except Exception as e:
             log.exception("Exception during extraction")
             log.debug("", exc_info=e)

tests/test_backup.py

@@ -0,0 +1,74 @@
+import gzip
+import hashlib
+from pathlib import Path
+from typing import IO
+from unittest.mock import patch
+
+import pytest
+
+from dissect.hypervisor.tools.backup import main
+
+
+@pytest.mark.parametrize(
+    "filename, expected",
+    [
+        (
+            "test9.vbk.gz",
+            [
+                (
+                    "6745a759-2205-4cd2-b172-8ec8f7e60ef8 (78a5467d-87f5-8540-9a84-7569ae2849ad_2d1bb20f-49c1-485d-a689-696693713a5a)/DEV__dev_nvme1n1",
+                    "337350cac29d2ed34c23ce9fc675950badf85fd2b694791abe6999d36f0dc1b3",
+                ),
+                (
+                    "6745a759-2205-4cd2-b172-8ec8f7e60ef8 (78a5467d-87f5-8540-9a84-7569ae2849ad_2d1bb20f-49c1-485d-a689-696693713a5a)/summary.xml",
+                    "d2b8f4d08e57a44b817b57d9c03e670c292e5a21e91fb5895b51e923781175e8",
+                ),
+            ],
+        ),
+        (
+            "test13.vbk.gz",
+            [
+                (
+                    "6745a759-2205-4cd2-b172-8ec8f7e60ef8 (3c834d56-37ac-8bd3-b946-30113c55c4b5)/BackupComponents.xml",
+                    "a9615b1cbce437074235ac194681d42e1f018f1c804277293dc24d4dd90eb504",
+                ),
+                (
+                    "6745a759-2205-4cd2-b172-8ec8f7e60ef8 (3c834d56-37ac-8bd3-b946-30113c55c4b5)/digest_47d9f323-442b-433d-bd4f-1ecb3fa97351",
+                    "d6f9dced7c58628a4648e1a5ed349609f11cefb0ac6721c35c5f943ac18aaf10",
+                ),
+                (
+                    "6745a759-2205-4cd2-b172-8ec8f7e60ef8 (3c834d56-37ac-8bd3-b946-30113c55c4b5)/GuestMembers.xml",
+                    "18228ae41c1e7ddb23ee6cfe49c9e2c1cdfe99393b45f20d7fbbdb5d247938b4",
+                ),
+                (
+                    "6745a759-2205-4cd2-b172-8ec8f7e60ef8 (3c834d56-37ac-8bd3-b946-30113c55c4b5)/summary.xml",
+                    "c93e3460a4495a96047fd8c1c80c3169782207e106cc3b5f5d9b5edae68eb3d9",
+                ),
+                (
+                    "6745a759-2205-4cd2-b172-8ec8f7e60ef8 (3c834d56-37ac-8bd3-b946-30113c55c4b5)/8b14f74c-360d-4d7a-98f7-7f4c5e737eb7",
+                    "e9ed281cf9c2fe1745e4eb9c926c1a64bd47569c48be511c5fdf6fd5793e5a77",
+                ),
+            ],
+        ),
+    ],
+)
+def test_backup_tool(
+    filename: str, expected: list[tuple[str, str]], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    Path_open = Path.open
+
+    def mock_open(self, mode: str, *args, **kwargs) -> IO:
+        if filename.endswith(".gz") and self.name == filename:
+            return gzip.open(self, mode)
+
+        return Path_open(self, mode, *args, **kwargs)
+
+    with patch.object(Path, "open", mock_open):
+        with monkeypatch.context() as m:
+            m.setattr("sys.argv", ["backup-extract", f"tests/data/{filename}", "-o", str(tmp_path)])
+            main()
+
+    for name, digest in expected:
+        out_path = tmp_path / name
+        assert out_path.exists()
+        assert hashlib.sha256(out_path.read_bytes()).hexdigest() == digest