add linuxnamespaces.tarutils

Move the generic tar utilities from the chroottar.py example into a
linuxnamespaces module as dealing with tar archives is a fairly common
thing when dealing with namespaces.

1 files changed, 12 insertions, 69 deletions

diff --git a/examples/chroottar.py b/examples/chroottar.py
index b210649..f43add4 100755
--- a/examples/chroottar.py
+++ b/examples/chroottar.py
@@ -11,78 +11,13 @@ import os
 import pathlib
 import socket
 import sys
-import tarfile
 import tempfile
-import typing
 
 if __file__.split("/")[-2:-1] == ["examples"]:
     sys.path.insert(0, "/".join(__file__.split("/")[:-2]))
 
 import linuxnamespaces
-
-
-class TarFile(tarfile.TarFile):
-    """Subclass of tarfile.TarFile that can read zstd compressed archives."""
-
-    OPEN_METH = {"zst": "zstopen"} | tarfile.TarFile.OPEN_METH
-
-    @classmethod
-    def zstopen(
-        cls,
-        name: str,
-        mode: typing.Literal["r", "w", "x"] = "r",
-        fileobj: typing.BinaryIO | None = None,
-    ) -> tarfile.TarFile:
-        if mode not in ("r", "w", "x"):
-            raise ValueError("mode must be 'r', 'w' or 'x'")
-        openobj: str | typing.BinaryIO = name if fileobj is None else fileobj
-        try:
-            import zstandard
-        except ImportError as err:
-            raise tarfile.CompressionError(
-                "zstandard module not available"
-            ) from err
-        if mode == "r":
-            zfobj = zstandard.open(openobj, "rb")
-        else:
-            zfobj = zstandard.open(
-                openobj,
-                mode + "b",
-                cctx=zstandard.ZstdCompressor(write_checksum=True, threads=-1),
-            )
-        try:
-            tarobj = cls.taropen(name, mode, zfobj)
-        except (OSError, EOFError, zstandard.ZstdError) as exc:
-            zfobj.close()
-            if mode == "r":
-                raise tarfile.ReadError("not a zst file") from exc
-            raise
-        except:
-            zfobj.close()
-            raise
-        # Setting the _extfileobj attribute is important to signal a need to
-        # close this object and thus flush the compressed stream.
-        # Unfortunately, tarfile.pyi doesn't know about it.
-        tarobj._extfileobj = False  # type: ignore
-        return tarobj
-
-    def get_comptype(self) -> str:
-        """Return the compression type used to compress the opened TarFile."""
-        # The tarfile module does not expose the compression method selected
-        # for open mode "r:*" in any way. We can guess it from the module that
-        # implements the fileobj.
-        compmodule = self.fileobj.__class__.__module__
-        try:
-            return {
-                "bz2": "bz2",
-                "gzip": "gz",
-                "lzma": "xz",
-                "_io": "tar",
-                "zstd": "zst",
-            }[compmodule]
-        except KeyError:
-            # pylint: disable=raise-missing-from  # no value in chaining
-            raise ValueError(f"cannot guess comptype for module {compmodule}")
+import linuxnamespaces.tarutils
 
 
 def main() -> None:
@@ -114,13 +49,19 @@ def main() -> None:
             parentsock.close()
             # Once we drop privileges via setreuid and friends, we may become
             # unable to open basetar or to chdir to tdir, so do those early.
-            with TarFile.open(args.basetar, "r:*") as tarf:
+            with linuxnamespaces.tarutils.ZstdTarFile.open(
+                args.basetar, "r:*"
+            ) as tarf:
                 os.chdir(tdir)
                 linuxnamespaces.unshare(
                     linuxnamespaces.CloneFlags.NEWUSER
                     | linuxnamespaces.CloneFlags.NEWNS
                 )
-                childsock.send(tarf.get_comptype().encode("ascii") + b"\0")
+                childsock.send(
+                    linuxnamespaces.tarutils.get_comptype(
+                        tarf
+                    ).encode("ascii") + b"\0",
+                )
                 childsock.recv(1)
                 childsock.close()
                 # The other process will now have set up our id mapping and
@@ -171,7 +112,9 @@ def main() -> None:
         if args.save and ret == 0:
             tmptar = f"{args.basetar}.new"
             try:
-                with TarFile.open(tmptar, "x:" + comptype) as tout:
+                with linuxnamespaces.tarutils.ZstdTarFile.open(
+                    tmptar, "x:" + comptype
+                ) as tout:
                     tout.add(tdir, ".")
                 os.rename(tmptar, args.basetar)
             except: