diff options
| -rwxr-xr-x | examples/chroottar.py | 81 | ||||
| -rw-r--r-- | linuxnamespaces/tarutils.py | 77 | ||||
| -rw-r--r-- | pyproject.toml | 3 |
3 files changed, 91 insertions, 70 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: diff --git a/linuxnamespaces/tarutils.py b/linuxnamespaces/tarutils.py new file mode 100644 index 0000000..c7a065c --- /dev/null +++ b/linuxnamespaces/tarutils.py @@ -0,0 +1,77 @@ +#!/usr/bin/python3 +# Copyright 2024 Helmut Grohne <helmut@subdivi.de> +# SPDX-License-Identifier: GPL-3 + +"""Extensions to the tarfile module. + * ZstdTarFile extends TarFile to deal with zstd-compressed archives. + * get_comptype guesses the compression used for an open TarFile. +""" + +import tarfile +import typing + + +class ZstdTarFile(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, + **kwargs: typing.Any, + ) -> 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, **kwargs) + 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(tarobj: tarfile.TarFile) -> str: + """Return the compression type used to compress the given 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 = tarobj.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}") diff --git a/pyproject.toml b/pyproject.toml index 481593a..04d6a5f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,8 +15,9 @@ requires-python = ">=3.9" # linuxnamespaces.systemd needs jeepney or dbussy, not both. jeepney = ["jeepney"] dbussy = ["dbussy"] +# linuxnamespaces.tarutils.ZstdTarFile +zstandard = ["zstandard"] test = ["pytest", "pytest-forked", "pytest-subtests"] -examples = ["zstandard"] [tool.black] line-length = 79 |