Add vmtar tool (#51)

Author: Schamper

.gitignore

@@ -6,6 +6,6 @@ dist/
 *.pyc
 __pycache__/
 .pytest_cache/
-tests/docs/api
-tests/docs/build
+tests/_docs/api
+tests/_docs/build
 .tox/

dissect/hypervisor/tools/vmtar.py

@@ -0,0 +1,20 @@
+import tarfile
+
+from dissect.hypervisor.util import vmtar
+
+
+def main() -> None:
+    # We just want to run the main function of the tarfile module, but with our VisorTarFile and is_tarfile functions
+    type(tarfile.main)(
+        tarfile.main.__code__,
+        tarfile.main.__globals__
+        | {
+            "TarFile": vmtar.VisorTarFile,
+            "is_tarfile": vmtar.is_tarfile,
+            "open": vmtar.open,
+        },
+    )()
+
+
+if __name__ == "__main__":
+    main()

dissect/hypervisor/util/vmtar.py

@@ -5,6 +5,8 @@
 
 import struct
 import tarfile
+from io import BytesIO
+from typing import BinaryIO, Final
 
 
 class VisorTarInfo(tarfile.TarInfo):
@@ -49,9 +51,72 @@ def _proc_member(self, tarfile: tarfile.TarFile) -> VisorTarInfo | tarfile.TarIn
         return super()._proc_member(tarfile)
 
 
-def VisorTarFile(*args, **kwargs) -> tarfile.TarFile:
-    return tarfile.TarFile(*args, **kwargs, tarinfo=VisorTarInfo)
+class VisorTarFile(tarfile.TarFile):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs, tarinfo=VisorTarInfo)
 
-
-def open(*args, **kwargs) -> tarfile.TarFile:
-    return tarfile.open(*args, **kwargs, tarinfo=VisorTarInfo)
+    @classmethod
+    def visoropen(cls, name: str, mode: str = "r", fileobj: BinaryIO | None = None, **kwargs) -> VisorTarFile:
+        """Open a visor tar file for reading. Supports gzip and lzma compression."""
+        if mode not in ("r",):
+            raise tarfile.TarError("visor currently only supports read mode")
+
+        try:
+            from gzip import GzipFile
+        except ImportError:
+            raise tarfile.CompressionError("gzip module is not available") from None
+
+        try:
+            from lzma import LZMAError, LZMAFile
+        except ImportError:
+            raise tarfile.CompressionError("lzma module is not available") from None
+
+        compressed = False
+
+        try:
+            t = cls.taropen(name, mode, fileobj, **kwargs)
+        except Exception:
+            try:
+                fileobj = GzipFile(name, mode + "b", fileobj=fileobj)
+            except OSError as e:
+                if fileobj is not None and mode == "r":
+                    raise tarfile.ReadError("not a visor file") from e
+                raise
+
+            try:
+                t = cls.taropen(name, mode, fileobj, **kwargs)
+            except Exception:
+                fileobj.seek(0)
+                fileobj = LZMAFile(fileobj or name, mode)  # noqa: SIM115
+
+                try:
+                    t = cls.taropen(name, mode, fileobj, **kwargs)
+                except (LZMAError, EOFError, OSError) as e:
+                    fileobj.close()
+                    if mode == "r":
+                        raise tarfile.ReadError("not a visor file") from e
+                    raise
+                except:
+                    fileobj.close()
+                    raise
+
+            compressed = True
+
+        # If we get here, we have a valid visor tar file
+        if fileobj is not None and compressed:
+            # Just read the entire file into memory, it's probably small
+            fileobj.seek(0)
+            fileobj = BytesIO(fileobj.read())
+
+        t = cls.taropen(name, mode, fileobj, **kwargs)
+
+        t._extfileobj = False
+        return t
+
+    # Only allow opening visor tar files
+    OPEN_METH: Final[dict[str, str]] = {"visor": "visoropen"}
+
+
+open = VisorTarFile.open
+
+is_tarfile = type(tarfile.is_tarfile)(tarfile.is_tarfile.__code__, tarfile.is_tarfile.__globals__ | {"open": open})

pyproject.toml

@@ -48,10 +48,11 @@ dev = [
 
 [project.scripts]
 envelope-decrypt = "dissect.hypervisor.tools.envelope:main"
+vmtar = "dissect.hypervisor.tools.vmtar:main"
 
 [tool.ruff]
 line-length = 120
-required-version = ">=0.9.0"
+required-version = ">=0.11.0"
 
 [tool.ruff.format]
 docstring-code-format = true
@@ -94,7 +95,7 @@ select = [
 ignore = ["E203", "B904", "UP024", "ANN002", "ANN003", "ANN204", "ANN401", "SIM105", "TRY003"]
 
 [tool.ruff.lint.per-file-ignores]
-"tests/docs/**" = ["INP001"]
+"tests/_docs/**" = ["INP001"]
 
 [tool.ruff.lint.isort]
 known-first-party = ["dissect.hypervisor"]