459 lines
15 KiB
Python
Executable File
459 lines
15 KiB
Python
Executable File
#!/usr/bin/env -S uv run --script
|
|
# /// script
|
|
# requires-python = ">=3.11"
|
|
# dependencies = ["httpx", "pillow"]
|
|
# ///
|
|
import argparse
|
|
import dataclasses
|
|
import logging
|
|
import os
|
|
import re
|
|
import struct
|
|
import subprocess
|
|
import sys
|
|
import urllib.parse
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
import httpx
|
|
from PIL import Image
|
|
|
|
run_command = partial(
|
|
subprocess.run,
|
|
check=True,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Image dimension reading (adapted from sort_images.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".gif"}
|
|
|
|
|
|
def read_dims_jpeg(image_path: Path) -> tuple[int, int]:
|
|
with image_path.open("rb") as f:
|
|
if f.read(2) != b"\xff\xd8":
|
|
raise ValueError(f"{image_path} is not a valid JPEG file")
|
|
f.seek(0)
|
|
try:
|
|
return Image.open(f).size
|
|
except Exception:
|
|
f.seek(0)
|
|
while True:
|
|
marker = f.read(1)
|
|
if not marker or marker != b"\xff":
|
|
raise ValueError(f"Invalid JPEG format in {image_path}")
|
|
marker_type = int.from_bytes(f.read(1), byteorder="big")
|
|
length = int.from_bytes(f.read(2), byteorder="big") - 2
|
|
is_sof = 0xC0 <= marker_type <= 0xCF and marker_type not in (
|
|
0xC4,
|
|
0xC8,
|
|
0xCC,
|
|
)
|
|
if is_sof:
|
|
f.seek(1, 1)
|
|
height = int.from_bytes(f.read(2), byteorder="big")
|
|
width = int.from_bytes(f.read(2), byteorder="big")
|
|
return width, height
|
|
f.seek(length, 1)
|
|
|
|
|
|
def read_dims_png(image_path: Path) -> tuple[int, int]:
|
|
with image_path.open("rb") as f:
|
|
if f.read(8) != b"\x89PNG\r\n\x1a\n":
|
|
raise ValueError(f"{image_path} is not a valid PNG file")
|
|
try:
|
|
f.seek(0)
|
|
with Image.open(f) as img:
|
|
return img.size
|
|
except Exception:
|
|
f.seek(12)
|
|
if f.read(4) != b"IHDR":
|
|
raise ValueError(f"IHDR chunk not found in {image_path}")
|
|
dims = f.read(8)
|
|
if len(dims) < 8:
|
|
raise ValueError(f"Truncated IHDR in {image_path}")
|
|
return struct.unpack(">II", dims)
|
|
|
|
|
|
def read_dims_webp(image_path: Path) -> tuple[int, int]:
|
|
with image_path.open("rb") as f:
|
|
header = f.read(12)
|
|
if len(header) < 12 or header[:4] != b"RIFF" or header[8:12] != b"WEBP":
|
|
raise ValueError(f"Not a valid WebP file: {image_path}")
|
|
while True:
|
|
chunk_header = f.read(8)
|
|
if len(chunk_header) < 8:
|
|
break
|
|
tag, length = struct.unpack("<4sI", chunk_header)
|
|
if tag == b"VP8X":
|
|
data = f.read(10)
|
|
return (int.from_bytes(data[4:7], "little") & 0xFFFFFF) + 1, (
|
|
int.from_bytes(data[7:10], "little") & 0xFFFFFF
|
|
) + 1
|
|
elif tag == b"VP8L":
|
|
data = f.read(5)
|
|
if data[0] != 0x2F:
|
|
raise ValueError("Invalid VP8L signature")
|
|
bits = struct.unpack("<I", data[1:5])[0]
|
|
return (bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1
|
|
elif tag == b"VP8 ":
|
|
data = f.read(10)
|
|
if data[3:6] != b"\x9d\x01\x2a":
|
|
raise ValueError("Invalid VP8 sync code")
|
|
w, h = struct.unpack("<HH", data[6:10])
|
|
return w & 0x3FFF, h & 0x3FFF
|
|
f.seek((length + 1) & ~1, 1)
|
|
raise ValueError(f"Could not find dimension chunks in {image_path}")
|
|
|
|
|
|
def read_dims(image_path: Path) -> tuple[int, int]:
|
|
match image_path.suffix.lower():
|
|
case ".jpg" | ".jpeg":
|
|
return read_dims_jpeg(image_path)
|
|
case ".png":
|
|
return read_dims_png(image_path)
|
|
case ".webp":
|
|
return read_dims_webp(image_path)
|
|
case _:
|
|
raise ValueError(f"Unsupported image format: {image_path.suffix}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def resolve_download_url_with_alldebrid(token: str, url: str) -> str:
|
|
resp = httpx.get(
|
|
"https://api.alldebrid.com/v4/link/unlock",
|
|
params={"link": url},
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
timeout=10,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
if data["status"] == "error":
|
|
raise ValueError(f"{data['error']['code']}: {data['error']['message']}")
|
|
return data["data"]["link"]
|
|
|
|
|
|
_CT_TO_EXT = {
|
|
"application/zip": ".zip",
|
|
"application/x-zip-compressed": ".zip",
|
|
"application/x-rar-compressed": ".rar",
|
|
"application/vnd.rar": ".rar",
|
|
"application/x-7z-compressed": ".7z",
|
|
"video/mp4": ".mp4",
|
|
"video/quicktime": ".mov",
|
|
"video/x-matroska": ".mkv",
|
|
}
|
|
|
|
_VIDEO_EXTS = {".mp4", ".mov", ".mkv"}
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class UrlMeta:
|
|
suggested_filename: str
|
|
suggested_extension: str
|
|
size_bytes: int
|
|
|
|
|
|
def _resolve_url(url: str) -> UrlMeta:
|
|
"""Send HEAD request; return filename, extension, and size from response headers."""
|
|
try:
|
|
resp = httpx.head(url, follow_redirects=True, timeout=15)
|
|
cd = resp.headers.get("content-disposition", "")
|
|
filename = None
|
|
if "filename=" in cd:
|
|
part = cd.split("filename=", 1)[1].strip().strip('"').strip("'")
|
|
filename = urllib.parse.unquote(part.split(";")[0].strip())
|
|
|
|
ct = resp.headers.get("content-type", "").split(";")[0].strip()
|
|
ext = _CT_TO_EXT.get(ct) or Path(urllib.parse.urlparse(url).path).suffix
|
|
|
|
if not filename:
|
|
filename = urllib.parse.unquote(Path(urllib.parse.urlparse(url).path).name)
|
|
if not Path(filename).suffix and ext:
|
|
filename = filename + ext
|
|
|
|
size = int(resp.headers.get("content-length", 0))
|
|
return UrlMeta(
|
|
suggested_filename=filename, suggested_extension=ext, size_bytes=size
|
|
)
|
|
except Exception:
|
|
url_path = Path(urllib.parse.urlparse(url).path)
|
|
name = urllib.parse.unquote(url_path.name)
|
|
return UrlMeta(
|
|
suggested_filename=name, suggested_extension=url_path.suffix, size_bytes=0
|
|
)
|
|
|
|
|
|
def _name_from_url(url: str) -> str:
|
|
parsed = urllib.parse.urlparse(url)
|
|
segment = parsed.path.rstrip("/").split("/")[-1]
|
|
return urllib.parse.unquote(segment) or parsed.netloc
|
|
|
|
|
|
def _has_images(d: Path) -> bool:
|
|
return any(f.suffix.lower() in IMAGE_EXTS for f in d.rglob("*") if f.is_file())
|
|
|
|
|
|
def _needs_resize(d: Path) -> bool:
|
|
images = [f for f in d.rglob("*") if f.is_file() and f.suffix.lower() in IMAGE_EXTS]
|
|
if not images:
|
|
return False
|
|
mean_mb = sum(f.stat().st_size for f in images) / len(images)
|
|
files_are_large = mean_mb >= 2 * 1_048_567
|
|
images_are_large = any(max(read_dims(img)) > 5000 for img in images)
|
|
|
|
return files_are_large or images_are_large
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Post-download: extract archives then resize
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _extract_and_process(output_dir: Path) -> None:
|
|
archives = [
|
|
f for pattern in ("*.zip", "*.rar", "*.7z") for f in output_dir.glob(pattern)
|
|
]
|
|
for archive in archives:
|
|
print(f"extracting {archive.name} …")
|
|
run_command(["unar", "-o", str(output_dir), str(archive)])
|
|
move_to_trash([archive])
|
|
|
|
subdirs = [d for d in output_dir.iterdir() if d.is_dir() and d.name != "resized"]
|
|
targets = subdirs if subdirs else [output_dir]
|
|
|
|
to_process = [d for d in targets if _needs_resize(d)]
|
|
if not to_process:
|
|
print("no dirs need resizing")
|
|
else:
|
|
resized_dir = output_dir / "resized"
|
|
resized_dir.mkdir(exist_ok=True)
|
|
originals = [
|
|
f
|
|
for d in to_process
|
|
for f in d.rglob("*")
|
|
if f.is_file() and f.suffix.lower() in IMAGE_EXTS
|
|
]
|
|
print(f"resizing {len(originals)} image(s) → {resized_dir}")
|
|
for img in originals:
|
|
run_command(
|
|
[
|
|
"imgz",
|
|
"resize",
|
|
"-o",
|
|
str(resized_dir / f"{img.stem}.resized.jpg"),
|
|
str(img),
|
|
]
|
|
)
|
|
if originals:
|
|
print(f"trashing {len(originals)} original(s)")
|
|
move_to_trash(originals)
|
|
rename_children(output_dir, flatten=True)
|
|
remove_empty_dirs(output_dir)
|
|
|
|
|
|
def remove_empty_dirs(root_dir: Path) -> None:
|
|
removed = [
|
|
d
|
|
for d in sorted(root_dir.rglob("*"), key=lambda p: len(p.parts), reverse=True)
|
|
if d.is_dir() and not any(d.iterdir())
|
|
]
|
|
for d in removed:
|
|
d.rmdir()
|
|
if removed:
|
|
print(f"removed {len(removed)} empty dir(s) in {root_dir}")
|
|
|
|
|
|
def move_to_trash(file_paths: list[Path]) -> None:
|
|
"""Move files to Trash using trash CLI (single batched call)."""
|
|
resolved: list[str] = []
|
|
for file_path in file_paths:
|
|
if not file_path.is_file():
|
|
raise FileNotFoundError(f"File not found: {file_path}")
|
|
resolved.append(str(file_path.resolve()))
|
|
|
|
if resolved:
|
|
run_command(["trash", *resolved])
|
|
|
|
|
|
def rename_children(
|
|
source_dir: Path,
|
|
flatten: bool,
|
|
filter_fn: Callable[[Path], bool] | None = None,
|
|
) -> None:
|
|
if not filter_fn:
|
|
filter_fn = is_image
|
|
|
|
def natural_sort_key(path: Path) -> tuple[str, ...]:
|
|
s = path.stem
|
|
return tuple(
|
|
f"{int(p):010d}" if p.isnumeric() else p
|
|
for p in re.findall(r"(\D+|\d+)", s)
|
|
)
|
|
|
|
assert source_dir.is_dir()
|
|
|
|
if flatten:
|
|
files = [f for f in source_dir.rglob("*") if f.is_file()]
|
|
else:
|
|
files = [f for f in source_dir.iterdir() if f.is_file()]
|
|
|
|
files = [f for f in files if filter_fn(f)]
|
|
|
|
files = sorted(files, key=natural_sort_key)
|
|
|
|
for i, f in enumerate(files):
|
|
ext = ".jpg" if f.suffix.lower() in {".jpeg", ".jpg"} else f.suffix
|
|
target_dir = f.parent
|
|
if flatten:
|
|
target_dir = source_dir
|
|
new_stem = source_dir.name
|
|
target = target_dir / f"{new_stem}__{i:04d}{ext}"
|
|
if target.is_file():
|
|
logging.error(f"target already exists: {target}")
|
|
continue
|
|
f.rename(target)
|
|
|
|
|
|
def is_image(f: Path) -> bool:
|
|
return f.is_file() and f.suffix.lower() in IMAGE_EXTS
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Subcommands
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def cmd_gallery(urls: list[str], name: str | None, cwd: Path | None = None) -> None:
|
|
resolved_name = name or _name_from_url(urls[0])
|
|
output_dir = (cwd or Path(".")) / resolved_name
|
|
|
|
if output_dir.exists() and _has_images(output_dir):
|
|
print(f"skipping gallery-dl: {output_dir} already has images")
|
|
else:
|
|
print(f"gallery-dl → {output_dir}")
|
|
args = ["uvx", "gallery-dl"]
|
|
if name or cwd:
|
|
args += ["-D", str(output_dir)]
|
|
run_command(args + urls)
|
|
|
|
_extract_and_process(output_dir)
|
|
|
|
|
|
def cmd_filehost(
|
|
urls: list[str], name: str | None, token: str, cwd: Path | None = None
|
|
) -> None:
|
|
resolved_name = name or _name_from_url(urls[0])
|
|
output_dir = (cwd or Path(".")) / resolved_name
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
for url in urls:
|
|
print(f"unrestricting {url}")
|
|
dl_url = resolve_download_url_with_alldebrid(token, url)
|
|
meta = _resolve_url(dl_url)
|
|
if meta.suggested_extension in _VIDEO_EXTS:
|
|
video_path = (cwd or Path(".")) / (resolved_name + meta.suggested_extension)
|
|
print(f"video detected → {video_path}")
|
|
if video_path.exists() and video_path.stat().st_size > 5 * 1024 * 1024:
|
|
print(f"skipping: {video_path.name} already downloaded")
|
|
continue
|
|
run_command(["aria2c", "-o", str(video_path), dl_url])
|
|
continue
|
|
|
|
dest = output_dir / meta.suggested_filename
|
|
if dest.exists() and dest.stat().st_size > 5 * 1024 * 1024:
|
|
print(f"skipping: {dest.name} already downloaded")
|
|
continue
|
|
print(f"downloading {meta.suggested_filename} → {output_dir}")
|
|
run_command(["aria2c", f"--dir={output_dir}", dl_url])
|
|
|
|
_extract_and_process(output_dir)
|
|
|
|
|
|
def cmd_process(dirs: list[str]) -> None:
|
|
for d in dirs:
|
|
p = Path(d)
|
|
if not p.exists():
|
|
print(f"skipping {p}: does not exist", file=sys.stderr)
|
|
continue
|
|
if _needs_resize(p):
|
|
resized_dir = p / "resized"
|
|
resized_dir.mkdir(exist_ok=True)
|
|
originals = [f for f in p.rglob("*") if is_image(f)]
|
|
print(f"resizing {len(originals)} image(s) in {p} → {resized_dir}")
|
|
for img in originals:
|
|
run_command(
|
|
[
|
|
"imgz",
|
|
"resize",
|
|
"-o",
|
|
str(resized_dir / f"{img.stem}.resized.jpg"),
|
|
str(img),
|
|
]
|
|
)
|
|
if originals:
|
|
print(f"trashing {len(originals)} original(s)")
|
|
move_to_trash(originals)
|
|
remove_empty_dirs(p)
|
|
else:
|
|
print(f"skipping {p}: already within size limits")
|
|
|
|
|
|
def cmd_download(
|
|
urls: list[str], name: str | None, token: str, cwd: Path | None = None
|
|
) -> None:
|
|
filehost_urls = [u for u in urls if "rg.to" in u or "rapidgator.net" in u]
|
|
gallery_urls = [u for u in urls if u not in set(filehost_urls)]
|
|
if filehost_urls:
|
|
cmd_filehost(filehost_urls, name, token, cwd)
|
|
if gallery_urls:
|
|
cmd_gallery(gallery_urls, name, cwd)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def main() -> None:
|
|
print(os.getenv("PATH"))
|
|
parser = argparse.ArgumentParser(prog="download_imageset.py")
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
|
|
gallery_p = sub.add_parser("gallery")
|
|
gallery_p.add_argument("urls", nargs="+")
|
|
gallery_p.add_argument("--name", default=None)
|
|
gallery_p.add_argument("--cwd", type=Path, default=None)
|
|
|
|
for cmd in ("filehost", "download"):
|
|
p = sub.add_parser(cmd)
|
|
p.add_argument("urls", nargs="+")
|
|
p.add_argument("--name", default=None)
|
|
p.add_argument("--cwd", default=None, type=Path)
|
|
p.add_argument("--token", default=os.getenv("ALLDEBRID_TOKEN"))
|
|
|
|
proc = sub.add_parser("process")
|
|
proc.add_argument("dirs", nargs="+")
|
|
|
|
args = parser.parse_args()
|
|
|
|
match args.command:
|
|
case "gallery":
|
|
cmd_gallery(args.urls, args.name, args.cwd)
|
|
case "filehost":
|
|
cmd_filehost(args.urls, args.name, args.token, args.cwd)
|
|
case "process":
|
|
cmd_process(args.dirs)
|
|
case "download":
|
|
cmd_download(args.urls, args.name, args.token, args.cwd)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|