playground/extract_archives.py

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)


def find_archive_files(root_dir: Path) -> list[Path]:
    extractables = []
    globs = ["*.rar", "*.zip", "*.7z", "*.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 re.search(f"\.zip|\.7z", archive.name):
        zip_part = re.search("(\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 "-p-",
        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 ValueError("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 ValueError("invalid password")

        raise


def extract_archive_7z(archive: Path, output_dir: Path, password: str | None = None):
    args = [
        "7z",
        "e",
        f'-p{password or ""}',
        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 FileNotFoundError("missing parts")
        if e.returncode == 2:
            raise ValueError("invalid password")

        raise


def extract_archive(archive_path: Path, output_dir: Path, passwords: list[str]) -> None:
    for p in ["", *passwords]:
        try:
            if ".rar" in archive_path.name:
                extract_archive_rar(archive_path, output_dir, p)
            else:
                extract_archive_7z(archive_path, output_dir, p)
        except Exception as e:
            logger.error(f"unhandled error: {e}")
            raise e


def delete_archive(archive: Path):
    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 parse_args():
    arger = argparse.ArgumentParser()
    arger.add_argument("archive_paths", type=Path, nargs="+", help="List of archive paths")
    arger.add_argument("--cwd", type=Path, default=Path.cwd(), help="working directory")
    arger.add_argument("--clean", type=bool, 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():
    args = parse_args()
    cwd: Path = args.cwd

    archives = find_archives(cwd)
    logger.info(f"found {len(archives)} archives")
    for f in archives:
        logger.info(f"\t{f.relative_to(cwd)}")

    typer.confirm("continue", default=True, abort=True)

    failed = []
    for i, f in enumerate(archives, start=1):
        progress = f"{i:02}/{len(archives):02}"
        logger.info(f"\nextracting {progress}: {f.relative_to(cwd)}")
        try:
            success = extract_archive(f, passwords=passwords)
            if success:
                if clean:
                    delete_archive(f)
                logger.info(f"done")
            else:
                logger.error("failed")
        except (ValueError, FileNotFoundError) as e:
            logger.error(str(e))
            failed.append(f)
    if failed:
        logger.info(f"\n\nfailed to extract {len(failed)} archives")
        for f in failed:
            logger.info(f"{f.relative_to(cwd)}")
    logger.info("\n\nfinished.")


if __name__ == "__main__":
    typer.run(main)