feat(extract-archives): add zip handling and find mode
This commit is contained in:
Regular → Executable
+94
-32
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user