diff --git a/extract_archives.py b/extract_archives.py old mode 100644 new mode 100755 index dd64e5e..29717e7 --- a/extract_archives.py +++ b/extract_archives.py @@ -1,3 +1,7 @@ +#!/usr/bin/env -S uv run +# /// script +# requires-python = ">=3.13" +# /// import argparse import logging import re @@ -21,9 +25,13 @@ dan4260 """.strip().splitlines(keepends=False) +class CorruptArchiveError(Exception): + pass + + def find_archive_files(root_dir: Path) -> list[Path]: extractables = [] - globs = ["*.rar", "*.zip", "*.7z", "*.zip.*", "*.7z.*"] + globs = ["*.rar", "*.zip*", "*.7z*"] for glob in globs: for f in root_dir.rglob(glob): extractables.append(f) @@ -34,8 +42,10 @@ def find_archive_files(root_dir: Path) -> list[Path]: def is_extractable(archive: Path) -> bool: - if re.search(f"\.zip|\.7z", archive.name): - zip_part = re.search("(\d+)$", archive.suffix) + 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: @@ -56,7 +66,7 @@ def extract_archive_rar(archive: Path, output_dir: Path, password: str | None = args = [ "rar", "e", - f"-p{password}" if password else "-p-", + *([f"-p{password}"] if password else []), f"{archive.name}", f"{output_dir}", ] @@ -65,13 +75,13 @@ def extract_archive_rar(archive: Path, output_dir: Path, password: str | None = 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") + 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 ValueError("invalid password") + raise LookupError("invalid password") raise @@ -80,7 +90,7 @@ def extract_archive_7z(archive: Path, output_dir: Path, password: str | None = N args = [ "7z", "e", - f'-p{password or ""}', + *(["-p", password] if password else []), f"-o{output_dir}", "-y", f"{archive}", @@ -91,26 +101,60 @@ def extract_archive_7z(archive: Path, output_dir: Path, password: str | None = N 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 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: - extract_archive_rar(archive_path, output_dir, p) + fn = extract_archive_rar + elif ".zip" == archive_path.suffix: + fn = extract_archive_zip else: - extract_archive_7z(archive_path, output_dir, p) + 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"unhandled error: {e}") - raise 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): @@ -118,11 +162,23 @@ def delete_archive(archive: Path): 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=Path, nargs="+", help="List of archive paths") + 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("--clean", type=bool, default=False, help="delete file(s) after successful extraction", action="store_true") + 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() @@ -132,37 +188,43 @@ def parse_args(): def main(): + logging.basicConfig(level=logging.INFO) + 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)}") + if args.find: + archives = find_archive_files(cwd) + for f in archives: + print(str(f.resolve())) + return - typer.confirm("continue", default=True, abort=True) + 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"\nextracting {progress}: {f.relative_to(cwd)}") + logger.info(f"extracting {progress}: {f.relative_to(cwd)}") + output_dir = f.parent / f.stem.strip() try: - success = extract_archive(f, passwords=passwords) - if success: - if clean: - delete_archive(f) - logger.info(f"done") - else: - logger.error("failed") + 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"\n\nfailed to extract {len(failed)} archives") + logger.info(f"failed 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) + main()