575 lines
21 KiB
Python
Executable File
575 lines
21 KiB
Python
Executable File
#!/usr/bin/env -S uv run --script
|
|
# /// script
|
|
# requires-python = ">=3.11"
|
|
# dependencies = ["httpx", "pillow", "click"]
|
|
# ///
|
|
import dataclasses
|
|
import logging
|
|
import os
|
|
import re
|
|
import shlex
|
|
import shutil
|
|
import struct
|
|
import subprocess
|
|
import tempfile
|
|
import urllib.parse
|
|
import zipfile
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
import click
|
|
import httpx
|
|
from PIL import Image, ImageOps
|
|
|
|
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())
|
|
|
|
|
|
_TRICKLEBAR_APP = Path("/Applications/TrickleBar.app")
|
|
|
|
|
|
def download_url(url: str, destination: Path) -> None:
|
|
destination = destination.resolve()
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
if _TRICKLEBAR_APP.is_dir():
|
|
logging.info(f"downloading with TrickleBar → {destination}")
|
|
download_url_with_tricklebar(url=url, destination=destination)
|
|
else:
|
|
logging.info(f"downloading with aria2 → {destination}")
|
|
download_url_with_aria2(url=url, destination=destination)
|
|
|
|
|
|
def download_url_with_aria2(url: str, destination: Path) -> None:
|
|
run_command(args=[
|
|
"aria2c", "--continue=true", f"--dir={destination.parent}",
|
|
f"--out={destination.name}", url,
|
|
])
|
|
|
|
|
|
def download_url_with_tricklebar(url: str, destination: Path) -> None:
|
|
with tempfile.TemporaryDirectory(prefix="imageset-") as temporary_directory:
|
|
directory = Path(temporary_directory)
|
|
completion_pipe = directory / "completion"
|
|
callback = directory / "on_done.sh"
|
|
os.mkfifo(path=completion_pipe, mode=0o600)
|
|
callback.write_text(data=(
|
|
"#!/bin/sh\n"
|
|
f"[ -p {shlex.quote(str(completion_pipe))} ] || exit 0\n"
|
|
f"printf '%s\\0' \"$1\" \"$2\" > {shlex.quote(str(completion_pipe))}\n"
|
|
))
|
|
callback.chmod(mode=0o700)
|
|
query = urllib.parse.urlencode(query={
|
|
"on_done": str(callback), "url": url,
|
|
"dir": str(destination.parent), "out": destination.name, "continue": "true",
|
|
}, quote_via=urllib.parse.quote)
|
|
logging.info("waiting for TrickleBar; press Ctrl+C to stop waiting")
|
|
run_command(args=["open", f"tricklebar://add-download?{query}"])
|
|
with completion_pipe.open(mode="rb") as completion:
|
|
status, downloaded_path, _ = completion.read().split(b"\0", maxsplit=2)
|
|
if status != b"complete":
|
|
raise click.ClickException(
|
|
f"TrickleBar download failed: {os.fsdecode(status)}"
|
|
)
|
|
if not downloaded_path or not Path(os.fsdecode(downloaded_path)).is_file():
|
|
raise click.ClickException("TrickleBar completed without a downloaded file")
|
|
if Path(os.fsdecode(downloaded_path)).resolve() != destination.resolve():
|
|
raise click.ClickException(
|
|
f"TrickleBar downloaded to {os.fsdecode(downloaded_path)} instead of {destination}; "
|
|
"update TrickleBar to a version supporting dir/out URL options"
|
|
)
|
|
if destination.with_name(f"{destination.name}.aria2").exists():
|
|
raise click.ClickException(f"TrickleBar download is unfinished: {destination}")
|
|
|
|
|
|
def _needs_resize(d: Path) -> bool:
|
|
images = [f for f in d.rglob("*") if f.is_file() and f.suffix.lower() in IMAGE_EXTS and not f.is_relative_to(d / "resized")]
|
|
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 = False
|
|
for image in images:
|
|
try:
|
|
if max(read_dims(image)) > 5000:
|
|
images_are_large = True
|
|
break
|
|
except Exception as error:
|
|
logging.warning(f"could not read dimensions of {image}; trying resizer: {error}")
|
|
images_are_large = True
|
|
break
|
|
|
|
return files_are_large or images_are_large
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Post-download: extract archives then resize
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
ImageProcessor = Callable[..., None]
|
|
|
|
|
|
def detect_resizer() -> ImageProcessor:
|
|
vips_binary = shutil.which(cmd="vips")
|
|
if vips_binary:
|
|
return partial(resize_with_vips, vips_binary=vips_binary)
|
|
return resize_with_pillow
|
|
|
|
|
|
def resize_with_vips(image_path: Path, output_path: Path, vips_binary: str) -> None:
|
|
try:
|
|
subprocess.run(
|
|
args=[
|
|
vips_binary,
|
|
"thumbnail",
|
|
str(image_path),
|
|
f"{output_path}[Q=75]",
|
|
"5000",
|
|
"--height",
|
|
"5000",
|
|
"--size",
|
|
"down",
|
|
],
|
|
check=True,
|
|
capture_output=True,
|
|
)
|
|
except (OSError, subprocess.CalledProcessError) as error:
|
|
logging.warning(f"vips failed for {image_path}; using Pillow: {error}")
|
|
resize_with_pillow(image_path=image_path, output_path=output_path)
|
|
|
|
|
|
def resize_with_pillow(image_path: Path, output_path: Path) -> None:
|
|
with Image.open(fp=image_path) as image:
|
|
with ImageOps.exif_transpose(image=image) as oriented:
|
|
oriented.thumbnail(size=(5000, 5000), resample=Image.Resampling.LANCZOS)
|
|
with oriented.convert(mode="RGB") as resized:
|
|
resized.save(fp=output_path, format="JPEG", quality=75)
|
|
|
|
|
|
def process_imageset(path: Path, processor: ImageProcessor) -> None:
|
|
if not path.exists():
|
|
logging.warning(f"skipping {path}: does not exist")
|
|
return
|
|
if path.is_file() and path.suffix.lower() in {".zip", ".rar", ".7z"}:
|
|
output_dir = path.with_suffix("")
|
|
archives = [path]
|
|
elif path.is_dir():
|
|
output_dir = path
|
|
archives = [f for f in path.iterdir() if f.is_file() and f.suffix.lower() in {".zip", ".rar", ".7z"}]
|
|
else:
|
|
logging.warning(f"skipping {path}: expected a directory or archive")
|
|
return
|
|
if output_dir.exists() and not output_dir.is_dir():
|
|
logging.warning(f"skipping {path}: {output_dir} is not a directory")
|
|
return
|
|
|
|
for archive in archives:
|
|
if archive.with_name(f"{archive.name}.aria2").exists():
|
|
logging.warning(f"skipping {archive}: download is unfinished")
|
|
return
|
|
logging.info(f"extracting {archive} → {output_dir}")
|
|
if archive.suffix.lower() == ".zip":
|
|
try:
|
|
with zipfile.ZipFile(file=archive) as zipped:
|
|
zipped.extractall(path=output_dir)
|
|
except zipfile.BadZipFile:
|
|
logging.warning(f"skipping {archive}: invalid ZIP file")
|
|
return
|
|
else:
|
|
run_command(args=["unar", "--output-directory", str(output_dir), str(archive)])
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
failed_images: set[Path] = set()
|
|
if _needs_resize(output_dir):
|
|
resized_dir = output_dir / "resized"
|
|
resized_dir.mkdir(exist_ok=True)
|
|
originals = [f for f in output_dir.rglob("*") if is_image(f) and not f.is_relative_to(resized_dir)]
|
|
logging.info(f"resizing {len(originals)} image(s) → {resized_dir}")
|
|
resized_originals: list[Path] = []
|
|
for img in originals:
|
|
output_path = resized_dir / f"{img.stem}.resized.jpg"
|
|
if output_path.is_file() and output_path.stat().st_size > 0:
|
|
logging.info(f"skipping {img}: {output_path} already resized")
|
|
resized_originals.append(img)
|
|
continue
|
|
try:
|
|
processor(image_path=img, output_path=output_path)
|
|
except Exception as error:
|
|
logging.warning(f"skipping {img}: resizing failed: {error}")
|
|
failed_images.add(img)
|
|
output_path.unlink(missing_ok=True)
|
|
continue
|
|
resized_originals.append(img)
|
|
if resized_originals:
|
|
logging.info(f"trashing {len(resized_originals)} original(s)")
|
|
move_to_trash(file_paths=resized_originals)
|
|
else:
|
|
logging.info(f"skipping resizing {output_dir}: already within size limits")
|
|
|
|
rename_children(
|
|
source_dir=output_dir,
|
|
flatten=True,
|
|
filter_fn=lambda image: is_image(image) and image not in failed_images,
|
|
)
|
|
remove_empty_dirs(root_dir=output_dir)
|
|
if archives and not failed_images:
|
|
move_to_trash(file_paths=archives)
|
|
|
|
|
|
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:
|
|
logging.info(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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_DEFAULT_CWD = Path("~/Downloads/temp/_imagesets").expanduser()
|
|
|
|
|
|
@click.group()
|
|
def cli() -> None:
|
|
pass
|
|
|
|
|
|
@cli.command("gallery")
|
|
@click.argument("urls", nargs=-1, required=True)
|
|
@click.option("--name", default=None)
|
|
@click.option("--cwd", type=Path, default=None)
|
|
def gallery_cmd(urls: tuple[str, ...], name: str | None, cwd: Path | None) -> None:
|
|
urls = list(urls)
|
|
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):
|
|
logging.info(f"skipping gallery-dl: {output_dir} already has images")
|
|
else:
|
|
logging.info(f"gallery-dl → {output_dir}")
|
|
args = ["uvx", "gallery-dl"]
|
|
if name or cwd:
|
|
args += ["-D", str(output_dir)]
|
|
run_command(args + urls)
|
|
|
|
process_imageset(path=output_dir, processor=detect_resizer())
|
|
|
|
|
|
@cli.command("filehost")
|
|
@click.argument("urls", nargs=-1, required=True)
|
|
@click.option("--name", default=None)
|
|
@click.option("--cwd", type=Path, default=_DEFAULT_CWD)
|
|
@click.option("--token", default=os.getenv("ALLDEBRID_TOKEN"))
|
|
def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str) -> 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:
|
|
logging.info(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)
|
|
logging.info(f"video detected → {video_path}")
|
|
if video_path.with_name(f"{video_path.name}.aria2").exists():
|
|
logging.warning(f"unfinished download; resuming {video_path}")
|
|
if (
|
|
video_path.exists()
|
|
and not video_path.with_name(f"{video_path.name}.aria2").exists()
|
|
and video_path.stat().st_size > 5 * 1024 * 1024
|
|
):
|
|
logging.info(f"skipping: {video_path.name} already downloaded")
|
|
continue
|
|
download_url(url=dl_url, destination=video_path)
|
|
continue
|
|
|
|
dest = output_dir / meta.suggested_filename
|
|
if dest.with_name(f"{dest.name}.aria2").exists():
|
|
logging.warning(f"unfinished download; resuming {dest}")
|
|
if (
|
|
dest.exists()
|
|
and not dest.with_name(f"{dest.name}.aria2").exists()
|
|
and dest.stat().st_size > 5 * 1024 * 1024
|
|
):
|
|
logging.info(f"skipping: {dest.name} already downloaded")
|
|
continue
|
|
logging.info(f"downloading {meta.suggested_filename} → {output_dir}")
|
|
download_url(url=dl_url, destination=dest)
|
|
|
|
process_imageset(path=output_dir, processor=detect_resizer())
|
|
|
|
def is_filehost(url: str) -> bool:
|
|
kws = ["rg.to", "rapidgator.net", "trbt.cc", "katfile"]
|
|
return any(kw in url for kw in kws)
|
|
|
|
@cli.command("download")
|
|
@click.argument("urls", nargs=-1, required=True)
|
|
@click.option("--name", default=None)
|
|
@click.option("--cwd", type=Path, default=_DEFAULT_CWD)
|
|
@click.option("--token", default=os.getenv("ALLDEBRID_TOKEN"))
|
|
@click.pass_context
|
|
def download_cmd(
|
|
ctx: click.Context, urls: tuple[str, ...], name: str | None, cwd: Path, token: str
|
|
) -> None:
|
|
filehost_urls = [u for u in urls if is_filehost(u)]
|
|
gallery_urls = [u for u in urls if u not in set(filehost_urls)]
|
|
if filehost_urls:
|
|
ctx.invoke(filehost_cmd, urls=tuple(filehost_urls), name=name, cwd=cwd, token=token)
|
|
if gallery_urls:
|
|
ctx.invoke(gallery_cmd, urls=tuple(gallery_urls), name=name, cwd=cwd)
|
|
|
|
|
|
@cli.command("process")
|
|
@click.argument("paths", nargs=-1, required=True, type=click.Path(path_type=Path))
|
|
def process_cmd(paths: tuple[Path, ...]) -> None:
|
|
"""Process directories and archives, extracting archives into sibling directories."""
|
|
processor = detect_resizer()
|
|
for path in paths:
|
|
process_imageset(path=path, processor=processor)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
|
cli()
|