playground/extract_archives.py

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.13"
# ///
import argparse
import logging
import re
import subprocess
import sys
from pathlib import Path

logger = logging.getLogger(__name__)

KNOWN_PASSWORDS = """
torgo
Yuzuki
oron.com
MonikGonPs
koth
drdoom-psuzy
psuzy
suzy
idle
dan4260
""".strip().splitlines(keepends=False)


class CorruptArchiveError(Exception):
    pass


def find_archive_files(root_dir: Path) -> list[Path]:
    extractables = []
    globs = ["*.rar", "*.zip*", "*.7z*"]
    for glob in globs:
        for f in root_dir.rglob(glob):
            extractables.append(f)

    extractables = [f for f in extractables if is_extractable(f)]

    return sorted(extractables, key=str)


def is_extractable(archive: Path) -> bool:
    if ".corrupt" in archive.name:
        return False
    if re.search(r"\.zip|\.7z", archive.name):
        zip_part = re.search(r"(\d+)$", archive.suffix)
        if zip_part and int(zip_part.group(1)) == 1:
            return True
        if not zip_part:
            return True

    if archive.name.endswith(".rar"):
        rar_part = re.search(r"\.part(\d+)", archive.name)
        if not rar_part:
            return True
        if rar_part and int(rar_part.group(1)) == 1:
            if archive.with_name(archive.name.replace("part1", "part2")).is_file():
                return True

    return False


def extract_archive_rar(archive: Path, output_dir: Path, password: str | None = None) -> None:
    args = [
        "rar",
        "e",
        *([f"-p{password}"] if password else []),
        f"{archive.name}",
        f"{output_dir}",
    ]
    try:
        subprocess.run(args, cwd=archive.parent, capture_output=True, text=True, check=True)
    except subprocess.CalledProcessError as e:
        output: str = e.stderr or e.stdout
        if "corrupt file or wrong password" in output.lower():
            raise LookupError("corrupt or invalid password")
        elif e.returncode == 3:
            raise FileNotFoundError("missing parts")
        elif e.returncode == 10:
            raise ValueError("not a rar file")
        elif e.returncode == 11:
            raise LookupError("invalid password")

        raise


def extract_archive_7z(archive: Path, output_dir: Path, password: str | None = None):
    args = [
        "7z",
        "e",
        *(["-p", password] if password else []),
        f"-o{output_dir}",
        "-y",
        f"{archive}",
    ]

    try:
        subprocess.run(args, capture_output=True, text=True, check=True)
    except subprocess.CalledProcessError as e:
        output: str = e.stderr or e.stdout
        if "missing volume" in output.lower():
            raise CorruptArchiveError("missing parts")
        if "wrong password" in output.lower():
            raise LookupError("invalid password")
        raise


def extract_archive_zip(archive: Path, output_dir: Path, password: str | None = None):
    args = [
        "unzip",
        "-o",  # overwrite files without prompting
        *(["-P", password] if password else []),
        f"{archive}",
        "-d",
        f"{output_dir}",
    ]

    try:
        subprocess.run(args, capture_output=True, text=True, check=True)
    except subprocess.CalledProcessError as e:
        output: str = e.stderr or e.stdout
        if "zipfile corrupt" in output.lower():
            raise CorruptArchiveError("corrupt zip file")
        if "cannot find zipfile" in output.lower():
            raise FileNotFoundError("zip file not found")
        if "incorrect password" in output.lower() or "password required" in output.lower():
            raise LookupError("invalid password")
        raise


def extract_archive(archive_path: Path, output_dir: Path, passwords: list[str]) -> None:
    for p in ["", *passwords]:
        try:
            fn: callable
            if ".rar" in archive_path.name:
                fn = extract_archive_rar
            elif ".zip" == archive_path.suffix:
                fn = extract_archive_zip
            else:
                fn = extract_archive_7z
            fn(archive_path, output_dir, p)
        except subprocess.CalledProcessError as e:
            logger.error(f"error extracting {archive_path}: {e.stderr}")
            raise
        except LookupError as e:
            logger.info(f"invalid password: {p!r}. {e}")
            continue
        except Exception as e:
            logger.error(f"failed to extract: {e}")
            raise


def delete_archive(archive: Path):
    logger.info(f"deleting: {archive.name}")
    archive.unlink()
    stem = archive.stem.rsplit(".", maxsplit=1)[0]
    for f in archive.parent.rglob(f"*{archive.suffix}"):
        if archive.stem.startswith(stem):
            logger.info(f"deleting: {f.name}")
            f.unlink()


def unlink_if_empty(dir_path: Path) -> None:
    """Unlink the directory if it is empty."""
    try:
        if not any(dir_path.iterdir()):
            dir_path.rmdir()
            logger.info(f"deleted empty directory: {dir_path}")
    except OSError as e:
        logger.error(f"error deleting directory: {e}")
        raise


def parse_args():
    arger = argparse.ArgumentParser()
    arger.add_argument("archive_paths", type=lambda v: Path(v.strip()), nargs="+", help="List of archive paths")
    arger.add_argument("--cwd", type=Path, default=Path.cwd(), help="working directory")
    arger.add_argument("--find", default=False, help="find archives and print to stdout", action="store_true")
    arger.add_argument("--clean", default=False, help="delete file(s) after successful extraction", action="store_true")

    if len(sys.argv) < 2:
        arger.print_help()
        exit(1)

    return arger.parse_args()


def main():
    logging.basicConfig(level=logging.INFO)

    args = parse_args()
    cwd: Path = args.cwd

    if args.find:
        archives = find_archive_files(cwd)
        for f in archives:
            print(str(f.resolve()))
        return

    archives = args.archive_paths

    failed = []
    for i, f in enumerate(archives, start=1):
        f: Path
        progress = f"{i:02}/{len(archives):02}"
        logger.info(f"extracting {progress}: {f.relative_to(cwd)}")
        output_dir = f.parent / f.stem.strip()
        try:
            output_dir.mkdir(exist_ok=True, parents=True)
            extract_archive(f, output_dir=output_dir, passwords=KNOWN_PASSWORDS)
            if args.clean:
                delete_archive(f)
        except CorruptArchiveError as e:
            if ".corrupt" not in f.name:
                f.rename(f.with_stem(f"{f.stem}.corrupt"))
        except (ValueError, FileNotFoundError) as e:
            logger.error(str(e))
            failed.append(f)
        finally:
            unlink_if_empty(output_dir)
    if failed:
        logger.info(f"failed to extract {len(failed)} archives")
        for f in failed:
            logger.info(f"{f.relative_to(cwd)}")


if __name__ == "__main__":
    main()