feat(extract-archives): add zip handling and find mode

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):
+    if ".corrupt" in archive.name:
-        zip_part = re.search("(\d+)$", archive.suffix)
+        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")
+            raise CorruptArchiveError("missing parts")
-        if e.returncode == 2:
+        if "wrong password" in output.lower():
-            raise ValueError("invalid password")
+            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}")
+            logger.error(f"failed to extract: {e}")
-            raise 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)
+    if args.find:
-    logger.info(f"found {len(archives)} archives")
+        archives = find_archive_files(cwd)
-    for f in archives:
+        for f in archives:
-        logger.info(f"\t{f.relative_to(cwd)}")
+            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)
+            output_dir.mkdir(exist_ok=True, parents=True)
-            if success:
+            extract_archive(f, output_dir=output_dir, passwords=KNOWN_PASSWORDS)
-                if clean:
+            if args.clean:
-                    delete_archive(f)
+                delete_archive(f)
-                logger.info(f"done")
+        except CorruptArchiveError as e:
-            else:
+            if ".corrupt" not in f.name:
-                logger.error("failed")
+                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()