unpack

#!/usr/bin/env python3
# Author: Jan Larres <jan@majutsushi.net>
# License: MIT/X11
#
# pyright: basic
# ruff: noqa: S603, S607

import argparse
import bz2
import gzip
import lzma
import os
import shutil
import subprocess
import sys
from functools import partial
from io import BufferedIOBase
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Callable, List, Optional

MIMETYPES = {
    "application/vnd.comicbook+zip": "zip",
    "application/vnd.comicbook-rar": "rar",
    "application/x-7z-compressed": "7z",
    "application/x-bzip-compressed-tar": "bztar",
    "application/x-compressed-tar": "gztar",
    "application/x-lha": "lha",
    "application/x-tar": "tar",
    "application/x-xz-compressed-tar": "xztar",
}


class UnpackError(Exception):
    pass


def main(args: argparse.Namespace) -> int:
    register_formats()

    try:
        do_unpack(args.archive, Path.cwd())
    except Exception as e:
        if args.verbose:
            raise
        print(e)
        return 1

    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Unpack various archives")
    parser.add_argument("archive", type=Path, help="the archive to unpack")
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        default=False,
        help="increase output verbosity",
    )
    return parser.parse_args()


def register_formats() -> None:
    # Register some additional extensions for ZIP files
    shutil.register_unpack_format(
        "zip2",
        [".jar", ".war", ".ear", ".xpi", ".cbz", ".epub"],
        shutil._unpack_zipfile,  # type: ignore[attr-defined]  # noqa: SLF001
        description="Additional ZIP extensions",
    )

    # Builtin functions
    shutil.register_unpack_format("bz2", [".bz2"], partial(run_builtin, bz2.open))
    shutil.register_unpack_format("gz", [".gz", ".Z"], partial(run_builtin, gzip.open))
    shutil.register_unpack_format(
        "xz",
        [".xz", ".lzma"],
        partial(run_builtin, lzma.open),
    )

    def register_tools(
        name: str,
        extensions: List[str],
        tools: List[Callable[[str, str], Optional[str]]],
    ) -> None:
        shutil.register_unpack_format(name, extensions, partial(run_tool, tools))

    register_tools("7z", [".7z"], [run_7z])
    register_tools("arc", [".arc"], [run_unar, run_arc])
    register_tools("iso", [".iso"], [run_unar])
    register_tools("lha", [".lha", ".lzh"], [run_lha])
    register_tools("zoo", [".zoo"], [run_unar])
    register_tools(
        "rar",
        [".rar", ".cbr"],
        # Run generic 'unrar' last since the free version has encoding issues
        [run_unrar_nonfree, run_7z, run_unar, run_unrar],
    )

    # Formats that require more complex unpacking
    shutil.register_unpack_format("deb", [".deb"], unpack_deb)
    shutil.register_unpack_format("rpm", [".rpm"], unpack_rpm)


def do_unpack(archive: Path, output_path: Path) -> None:
    print(f"Unpacking {archive} ...")
    with TemporaryDirectory(prefix="unpack-") as tmpdir:
        try:
            shutil.unpack_archive(archive, tmpdir)
        except shutil.ReadError as e:
            # Try getting the mimetype with Gio
            from gi.repository import Gio  # type: ignore[import-untyped]

            with archive.open("rb") as f:
                mime = Gio.content_type_guess(
                    filename=os.fspath(archive), data=f.read(1024)
                )[0]
            if mime in MIMETYPES:
                try:
                    shutil.unpack_archive(archive, tmpdir, format=MIMETYPES[mime])
                except shutil.ReadError as e2:
                    raise UnpackError(str(e2)) from e
            else:
                raise UnpackError(str(e)) from e

        items = list(Path(tmpdir).iterdir())
        if len(items) == 0:
            # Didn't unpack anything?
            raise UnpackError("No files were unpacked")
        if len(items) == 1:
            # Move single items directly into the target directory
            item = items[0]
            shutil.move(os.fspath(item), os.fspath(output_path))
            final_path = output_path / item.name
            print(f"Unpacked {archive} to {get_display_path(final_path)}")
        else:
            # Move multiple items into a new directory named after the archive
            output_dirname = archive.stem
            if output_dirname.endswith(".tar"):
                output_dirname = output_dirname[: -len(".tar")]
            output_path = output_path / output_dirname
            if output_path.exists():
                raise UnpackError(f"Directory {output_path} already exists")
            output_path.mkdir()
            for item in items:
                shutil.move(os.fspath(item), os.fspath(output_path))
            print(f"Unpacked {archive} to {get_display_path(output_path)}")


def run_builtin(
    func: Callable[[str], BufferedIOBase],
    archive: str,
    output_dir: str,
) -> None:
    archive_path = Path(archive)
    with func(archive) as src, Path(output_dir, archive_path.stem).open("wb") as dst:
        shutil.copyfileobj(src, dst)


def run_tool(
    tool_funcs: List[Callable[[str, str], Optional[str]]], archive: str, output_dir: str
) -> None:
    tried = []
    for func in tool_funcs:
        missing_tool = func(archive, output_dir)
        if missing_tool is None:
            break
        tried.append(missing_tool)
    else:
        raise UnpackError(
            f"Unable to find a tool for unpacking {archive}\n"
            f"Supported tools: {','.join(tried)}"
        )


def run_7z(archive: str, output_dir: str) -> Optional[str]:
    tool = "7z"
    if shutil.which(tool):
        subprocess.run([tool, "x", f"-o{output_dir}", archive], check=True)
        return None
    return tool


def run_arc(archive: str, output_dir: str) -> Optional[str]:
    tool = "arc"
    if shutil.which(tool):
        subprocess.run([tool, "x", archive], cwd=output_dir, check=True)
        return None
    return tool


def run_lha(archive: str, output_dir: str) -> Optional[str]:
    tool = "lha"
    if shutil.which(tool):
        subprocess.run([tool, "-x", archive], cwd=output_dir, check=True)
        return None
    return tool


def run_unar(archive: str, output_dir: str) -> Optional[str]:
    tool = "unar"
    if shutil.which(tool):
        subprocess.run([tool, "-o", output_dir, archive], check=True)
        return None
    return tool


def run_unrar(archive: str, output_dir: str) -> Optional[str]:
    tool = "unrar"
    if shutil.which(tool):
        subprocess.run([tool, "x", archive, f"{output_dir}/"], check=True)
        return None
    return tool


def run_unrar_nonfree(archive: str, output_dir: str) -> Optional[str]:
    tool = "unrar-nonfree"
    if shutil.which(tool):
        subprocess.run([tool, "x", archive, f"{output_dir}/"], check=True)
        return None
    return tool


def unpack_deb(archive: str, output_dir: str) -> None:
    subprocess.run(["ar", "x", archive], cwd=output_dir, check=True)
    data = next(Path(output_dir).glob("data.tar.*"))
    do_unpack(data, Path(output_dir))
    data.unlink()
    control = next(Path(output_dir).glob("control.tar.*"))
    do_unpack(control, Path(output_dir))
    control.unlink()


def unpack_rpm(archive: str, output_dir: str) -> None:
    proc = subprocess.Popen(["rpm2cpio", archive], stdout=subprocess.PIPE)
    subprocess.run(["cpio", "-idmv"], cwd=output_dir, stdin=proc.stdout, check=True)
    proc.wait()


def get_display_path(abs_path: Path) -> Path:
    try:
        return abs_path.relative_to(Path.cwd())
    except ValueError:
        return abs_path


if __name__ == "__main__":
    sys.exit(main(parse_args()))