#!/usr/bin/python3
# Copyright 2024 Helmut Grohne <helmut@subdivi.de>
# SPDX-License-Identifier: LGPL-2.0-or-later

"""Extensions to the tarfile module.
 * ZstdTarFile extends TarFile to deal with zstd-compressed archives.
 * get_comptype guesses the compression used for an open TarFile.
 * XAttrTarFile extends TarFile to map extended attributes to PAX headers.
"""

import collections.abc
import contextlib
import copy
import os
import tarfile
import typing


if typing.TYPE_CHECKING:
    from typing import override
else:

    def override(method):
        """typing.override stub for Python < 3.12."""
        return method


TarPath = str | bytes | os.PathLike[str] | os.PathLike[bytes]


_zstd_exceptions: set[type[BaseException]] = set()


def _zstd_open_package(
    openobj: str | typing.BinaryIO,
    mode: typing.Literal["r", "w", "x"],
    compresslevel: int | None = None,
    threads: int | None = None,
) -> typing.BinaryIO:
    """Open a zstd-compressed file using the zstandard module."""
    import zstandard  # Propagate ImportError
    _zstd_exceptions.add(zstandard.ZstdError)
    if mode == "r":
        return zstandard.open(openobj, "rb")
    if threads is None:
        threads = -1
    if compresslevel is not None:
        if compresslevel > 22:
            raise ValueError(f"invalid compression level {compresslevel}")
        cctx = zstandard.ZstdCompressor(
            write_checksum=True, threads=threads, level=compresslevel
        )
    else:
        cctx = zstandard.ZstdCompressor(
            write_checksum=True, threads=threads
        )
    return zstandard.open(
        openobj,
        # Employ a cast to make mypy happy.
        typing.cast(typing.Literal["rb", "wb", "xb"], mode + "b"),
        cctx=cctx
    )


def _zstd_open_stdlib(
    openobj: str | typing.BinaryIO,
    mode: typing.Literal["r", "w", "x"],
    compresslevel: int | None = None,
) -> typing.BinaryIO:
    """Open a zstd-compressed file using the compression.zstd module since
    Python 3.14.
    """
    import compression.zstd  # Propagate ImportError
    _zstd_exceptions.add(compression.zstd.ZstdError)
    return compression.zstd.ZstdFile(
        openobj,
        # Employ a cast to make mypy happy.
        typing.cast(typing.Literal["rb", "wb", "xb"], mode + "b"),
        level=compresslevel,
    )


def _zstd_open(
    openobj: str | typing.BinaryIO,
    mode: typing.Literal["r", "w", "x"],
    compresslevel: int | None = None,
    threads: int | None = None,
) -> typing.BinaryIO:
    """Open a zstd-compressed fileobj using a suitable decompression module."""
    try_stdlib = True
    if threads == 1 or mode == "r":
        try:
            return _zstd_open_stdlib(openobj, mode, compresslevel)
        except ImportError:
            try_stdlib = False
    with contextlib.suppress(ImportError):
        return _zstd_open_package(openobj, mode, compresslevel, threads)
    if try_stdlib:
        with contextlib.suppress(ImportError):
            return _zstd_open_stdlib(openobj, mode, compresslevel)
    raise tarfile.CompressionError(
        "neither compression.zstd nor zstandard module available"
    )


class ZstdTarFile(tarfile.TarFile):
    """Subclass of tarfile.TarFile that can read zstd compressed archives."""

    # mypy types OPEN_METH as Mapping rather than dict while it really is a
    # dict. Hence, it complains that there is no __or__ for dict and Mapping.
    OPEN_METH = {
        "zst": "zstopen",
    } | tarfile.TarFile.OPEN_METH  # type: ignore[operator]

    @classmethod
    def zstopen(
        cls,
        name: str,
        mode: typing.Literal["r", "w", "x"] = "r",
        fileobj: typing.BinaryIO | None = None,
        *,
        compresslevel: int | None = None,
        threads: int | None = None,
        **kwargs: typing.Any,
    ) -> tarfile.TarFile:
        """Open a zstd compressed tar archive with the given name for reading
        or writing. Appending is not supported. The class allows customizing
        the compression level and the compression concurrency (default
        parallel) while decompression ignores those arguments.
        """
        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
        zfobj = _zstd_open(openobj, mode, compresslevel, threads)
        # Work around https://github.com/python/mypy/issues/21092
        exceptions_to_catch: tuple[type[BaseException], ...] = (
            OSError,
            EOFError,
            *_zstd_exceptions,
        )
        try:
            tarobj = cls.taropen(name, mode, zfobj, **kwargs)
        except exceptions_to_catch 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__
    compmodule = compmodule.removeprefix("compression.")  # Python >= 3.14
    compmodule = compmodule.split(".", 1)[0]  # Strip submodules
    try:
        return {
            "bz2": "bz2",
            "gzip": "gz",
            "lzma": "xz",
            "_io": "tar",
            "zstandard": "zst",
            "zstd": "zst",
        }[compmodule]
    except KeyError:
        # pylint: disable=raise-missing-from  # no value in chaining
        raise ValueError(f"cannot guess comptype for module {compmodule}")


class XAttrTarFile(tarfile.TarFile):
    """A subclass to tarfile.TarFile that adds support for extended attributes
    via SCHILY.xattr.* PAX headers to extraction and creation of archives. It
    can be used as a mixin class with others as it does not add any state.
    """

    @override
    def extract(
        self,
        member: tarfile.TarInfo | str,
        path: TarPath = "",
        set_attrs: bool = True,
        **kwargs: typing.Any,
    ) -> None:
        """Refer to tarfile.TarFile.extract. In addition, SCHILY.xattr.* PAX
        headers are examined and applied as extended attributes if set_attrs is
        true-ish.
        """
        if not set_attrs:
            super().extract(member, path, False, **kwargs)
            return

        # We also need the tarinfo, so mimic the start of the built-in extract.
        if isinstance(member, str):
            tarinfo = self.getmember(member)
        else:
            tarinfo = member

        # extract applies the possible filter functions, but does not yield
        # back the filtered tarinfo.
        try:
            get_filter_function = getattr(self, "_get_filter_function")
        except AttributeError:  # Filters introduced in 3.12
            super().extract(tarinfo, path, True, **kwargs)
        else:
            # What follows is a copy of the extract method implementation
            filter_function = get_filter_function(kwargs.pop("filter", None))
            tarinfo, _unfiltered = (
                self._get_extract_tarinfo(  # type: ignore[attr-defined]
                    member, filter_function, path
                )
            )
            if tarinfo is None:
                return
            self._extract_one(  # type: ignore[attr-defined]
                tarinfo, path, True, **kwargs
            )

        # mypy is unhappy about the next line, but we have the same code in
        # TarFile.extract and if it bails here, it also bails there.
        path = os.path.join(path, tarinfo.name)  # type: ignore

        for attr, value in tarinfo.pax_headers.items():
            if not attr.startswith("SCHILY.xattr."):
                continue
            attr = attr.removeprefix("SCHILY.xattr.")
            os.setxattr(
                path,
                attr,
                value.encode(self.encoding or "utf8", "surrogateescape"),
                follow_symlinks=False,
            )

    @override
    def gettarinfo(
        self,
        name: TarPath | None = None,
        arcname: str | None = None,
        fileobj: typing.IO[bytes] | None = None,
    ) -> tarfile.TarInfo:
        tarinfo = super().gettarinfo(name, arcname, fileobj)
        path: int | TarPath
        if fileobj is not None:
            path = fileobj.fileno()
        elif name is not None:
            path = name
        else:
            raise ValueError("gettarinfo requires a name or fileobj")
        dereference = True if self.dereference is None else self.dereference
        for attr in os.listxattr(path, follow_symlinks=dereference):
            key = "SCHILY.xattr." + attr
            value = os.getxattr(
                path, attr, follow_symlinks=dereference
            ).decode(self.encoding or "utf8", "surrogateescape")
            # TarInfo.pax_headers is designated as (read-only) Mapping, but it
            # really is a writable dict.
            tarinfo.pax_headers[key] = value  # type: ignore[index]
        return tarinfo


TarFilterFunction = collections.abc.Callable[
    [tarfile.TarInfo, str], tarfile.TarInfo | None
]


class _ComposeTarFilters:
    def __init__(self, filters: tuple[TarFilterFunction, ...]) -> None:
        self.filters = filters

    def __call__(
        self, member: tarfile.TarInfo, path: str
    ) -> tarfile.TarInfo | None:
        for ff in self.filters:
            result = ff(member, path)
            if result is None:
                return None
            member = result
        return member


def compose_filter(*filters: TarFilterFunction) -> TarFilterFunction:
    """Compose tarfile filters by chaining them together.

    The given filters are invoked in the given order. The chain terminates when
    a filter skips a member (returns None) or aborts extraction (raises and
    exception). The member object is mutated by each filter in order.
    """
    match filters:
        case ():
            return tarfile.fully_trusted_filter
        case (single_filter,):
            return single_filter
    return _ComposeTarFilters(filters)


def skipdevice_filter(
    member: tarfile.TarInfo, _path: str
) -> tarfile.TarInfo | None:
    """A tarfile extraction filter that skips over device files.

    Unlike the builtin data_filter, it does not exclude named pipes.
    """
    if member.isblk() or member.ischr():
        return None
    return member


def _relative_to_root(
    name: str,
    *,
    sep: str = os.path.sep,
    pardir: str = os.path.pardir,
    curdir: str = os.path.curdir,
) -> str:
    parts: list[str] = []
    for part in name.split(sep):
        if part == pardir:
            if parts:
                parts.pop()
        elif part not in ('', curdir):
            parts.append(part)
    if not parts:
        return curdir
    return sep.join(parts)


def sanitizepath_filter(
    member: tarfile.TarInfo, _path: str
) -> tarfile.TarInfo | None:
    """A tarfile extraction filter that normalizes paths such that all file
    system operations during extraction operate below the extraction root.
    """
    name = member.name
    new_name: str | None = _relative_to_root(name)
    if new_name in (name, name.removeprefix("./")):
        new_name = None
    # Handle hard links as they are resolved at extraction time, but do not
    # handle symlinks as they are resolved at runtime.
    new_linkname: str | None
    if member.islnk():
        linkname = member.linkname
        new_linkname = _relative_to_root(linkname)
        if new_linkname in (linkname, linkname.removeprefix("./")):
            new_linkname = None
    else:
        new_linkname = None
    if new_name is None and new_linkname is None:
        return member
    new = copy.copy(member)
    if new_name is not None:
        new.name = new_name
    if new_linkname is not None:
        new.linkname = new_linkname
    return new
