J621/backend/apps/library/services.py

import hashlib
import json
import logging
import mimetypes
import os
import re
import shutil
import subprocess
from pathlib import Path

import imagehash
from django.conf import settings
from django.core import signing
from django.http import FileResponse, Http404, HttpResponse
from django.utils.text import get_valid_filename
from PIL import Image

from .models import MediaItem, MediaLocation

logger = logging.getLogger(__name__)

HASH_FIELDS = ("ahash", "dhash", "phash", "whash")
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".apng", ".webp"}

ALLOWED_EXTENSIONS = {
    ".jpg",
    ".jpeg",
    ".png",
    ".gif",
    ".apng",
    ".webp",
    ".mp4",
    ".webm",
}
VIDEO_EXTENSIONS = {".mp4", ".webm"}
UPLOAD_FILE_SALT = "j621.upload-file"
MEDIA_FILE_SALT = "j621.media-file"
CHUNK_SIZE = 1024 * 1024
RANGE_RE = re.compile(r"bytes=(\d*)-(\d*)$")


def compute_md5(path):
    digest = hashlib.md5()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(CHUNK_SIZE), b""):
            digest.update(chunk)
    return digest.hexdigest()


def index_file(path, folder):
    """Index one file into MediaItem/MediaLocation.

    Returns (item, created_item, location, created_location).
    """
    path = Path(path)
    folder = Path(folder)
    stat = path.stat()
    md5 = compute_md5(path)
    item, created_item = MediaItem.objects.get_or_create(
        md5=md5, defaults={"size": stat.st_size}
    )
    if not created_item and item.size != stat.st_size:
        item.size = stat.st_size
        item.save(update_fields=["size", "updated_at"])
    location, created_location = MediaLocation.objects.update_or_create(
        item=item,
        path=str(path),
        defaults={
            "rel_path": str(path.relative_to(folder)),
            "mtime": stat.st_mtime,
        },
    )
    return item, created_item, location, created_location


def signed_media_url(item, user, action="raw", request=None):
    """Media URL that <img>/<video> tags can load for a signed-in user.

    With a ``request`` the URL is absolute, so the SPA also works when it is
    served from a different origin; without one it stays relative.
    """
    path = f"/api/files/J-{item.id}/{action}/"
    if user is not None and getattr(user, "is_authenticated", False):
        signature = signing.dumps(
            {"item": item.id, "user": user.id, "action": action},
            salt=MEDIA_FILE_SALT,
        )
        path = f"{path}?sig={signature}"
    if request is None:
        return path
    return request.build_absolute_uri(path)


def rename_location_to_j_id(item, location):
    """Name a freshly indexed copy J-<id>.<ext> inside its own folder."""
    path = Path(location.path)
    if not path.exists():
        return location
    if path.stem == f"J-{item.id}":
        return location
    target = unique_destination(path.parent, f"J-{item.id}{path.suffix}")
    path.rename(target)
    parent = Path(location.rel_path).parent
    location.path = str(target)
    location.rel_path = (
        str(parent / target.name) if str(parent) != "." else target.name
    )
    location.save(update_fields=["path", "rel_path"])
    return location


def compute_visual_hashes(path):
    """Perceptual hashes for an image file (empty dict for other files)."""
    path = Path(path)
    if path.suffix.lower() not in IMAGE_EXTENSIONS:
        return {}
    try:
        with Image.open(path) as image:
            converted = image.convert("RGB")
            return {
                "ahash": str(imagehash.average_hash(converted, hash_size=8)),
                "dhash": str(imagehash.dhash(converted, hash_size=8)),
                "phash": str(imagehash.phash(converted, hash_size=8)),
                "whash": str(imagehash.whash(converted, hash_size=8)),
            }
    except Exception:  # noqa: BLE001 - hashing must never break indexing
        logger.exception("Could not compute visual hashes for %s", path)
        return {}


def ensure_visual_hashes(item):
    """Fill in missing perceptual hashes for a media item."""
    if all(getattr(item, field) for field in HASH_FIELDS):
        return item
    location = item.locations.first()
    if location is None:
        return item
    hashes = compute_visual_hashes(location.path)
    if not hashes:
        return item
    for field, value in hashes.items():
        setattr(item, field, value)
    item.save(update_fields=[*hashes.keys(), "updated_at"])
    return item


def parse_tags(raw):
    """Normalize a comma-separated string or JSON list into a list of tags."""
    if raw is None:
        return None
    if isinstance(raw, (list, tuple)):
        values = raw
    else:
        text = str(raw).strip()
        if not text:
            return None
        if text.startswith("["):
            try:
                values = json.loads(text)
            except json.JSONDecodeError:
                values = text.split(",")
        else:
            values = text.split(",")
    cleaned = []
    for value in values:
        tag = str(value).strip()
        if tag and tag not in cleaned:
            cleaned.append(tag[:100])
    return cleaned


def unique_destination(folder, filename):
    """Return a non-existing path inside folder for the given filename."""
    folder = Path(folder)
    folder.mkdir(parents=True, exist_ok=True)
    filename = get_valid_filename(Path(filename).name) or "upload"
    candidate = folder / filename
    stem, suffix = candidate.stem, candidate.suffix
    counter = 1
    while candidate.exists():
        candidate = folder / f"{stem}-{counter}{suffix}"
        counter += 1
    return candidate


class RangeFileWrapper:
    """Iterate over a limited byte range of an open file."""

    def __init__(self, file, length, chunk_size=64 * 1024):
        self.file = file
        self.remaining = length
        self.chunk_size = chunk_size

    def __iter__(self):
        return self

    def __next__(self):
        if self.remaining <= 0:
            raise StopIteration
        data = self.file.read(min(self.chunk_size, self.remaining))
        if not data:
            raise StopIteration
        self.remaining -= len(data)
        return data

    def close(self):
        self.file.close()


def serve_file(request, path, download=False):
    """Serve a file with HTTP range support (needed for video seeking)."""
    path = Path(path)
    if not path.is_file():
        raise Http404
    size = path.stat().st_size
    content_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
    range_header = request.headers.get("Range", "").strip()
    if range_header:
        match = RANGE_RE.match(range_header)
        if match:
            start_raw, end_raw = match.groups()
            if start_raw == "" and end_raw:
                length = min(int(end_raw), size)
                start, end = size - length, size - 1
            else:
                start = int(start_raw or 0)
                end = min(int(end_raw) if end_raw else size - 1, size - 1)
            if start >= size or start > end:
                response = HttpResponse(status=416)
                response["Content-Range"] = f"bytes */{size}"
                return response
            length = end - start + 1
            handle = open(path, "rb")
            handle.seek(start)
            response = FileResponse(
                RangeFileWrapper(handle, length),
                status=206,
                content_type=content_type,
            )
            response["Content-Length"] = str(length)
            response["Content-Range"] = f"bytes {start}-{end}/{size}"
            response["Accept-Ranges"] = "bytes"
            return response
    response = FileResponse(
        open(path, "rb"),
        content_type=content_type,
        as_attachment=download,
        filename=path.name,
    )
    response["Accept-Ranges"] = "bytes"
    return response


class DownloadCancelled(Exception):
    """Raised when a streamed download is cancelled by the user."""


class RemoteUrlError(ValueError):
    """The URL is not an allowed e621 media URL."""


def validate_remote_url(url):
    """Only http(s) URLs on the known e621 media hosts may be fetched.

    Without this the download paths are an SSRF hole: any uploader could make
    the server fetch internal addresses (127.0.0.1, LAN services, cloud
    metadata) and read the response back through the library.
    """
    from urllib.parse import urlparse

    parsed = urlparse(str(url or "").strip())
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise RemoteUrlError("Only http(s) URLs can be fetched.")
    if parsed.hostname not in settings.E621_MEDIA_HOSTS:
        raise RemoteUrlError("That host is not an allowed e621 media host.")
    return url


def open_remote(url, *, max_redirects=3, **kwargs):
    """GET an allowlisted URL, re-validating every redirect hop.

    Returns a streaming ``requests`` response. Redirects are followed
    manually so a hop cannot jump to an internal host.
    """
    import requests
    from urllib.parse import urljoin

    current = url
    for _ in range(max_redirects + 1):
        validate_remote_url(current)
        response = requests.get(current, allow_redirects=False, **kwargs)
        if response.is_redirect or response.is_permanent_redirect:
            location = response.headers.get("Location")
            response.close()
            if not location:
                raise RemoteUrlError("The remote server redirected without a target.")
            current = urljoin(current, location)
            continue
        return response
    raise RemoteUrlError("Too many redirects from the remote server.")


def download_file(
    url,
    destination,
    progress_callback=None,
    should_cancel=None,
    read_timeout=60,
):
    """Stream a remote file into destination (used by Download to Library).

    The read timeout bounds how long a stalled connection can block the
    worker: without it a hung socket would keep a job "downloading" forever
    and the cancel flag could never be observed.
    """
    headers = {"User-Agent": settings.USER_AGENT}
    with open_remote(
        url, headers=headers, stream=True, timeout=(10, read_timeout)
    ) as response:
        response.raise_for_status()
        total = int(response.headers.get("content-length") or 0)
        downloaded = 0
        with open(destination, "wb") as handle:
            for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
                if should_cancel is not None and should_cancel():
                    raise DownloadCancelled("download cancelled")
                if chunk:
                    handle.write(chunk)
                    downloaded += len(chunk)
                    if progress_callback is not None:
                        progress_callback(downloaded, total)
        if progress_callback is not None:
            progress_callback(downloaded, total or downloaded)


E621_DESCRIPTION_LIMIT = 20000


def trim_e621_post(post):
    """Keep a compact, size-bounded copy of an e621 post payload."""
    if not isinstance(post, dict):
        return None
    file_data = post.get("file") or {}
    relationships = post.get("relationships") or {}
    score = post.get("score") or {}
    categories = post.get("tags") or {}
    return {
        "id": post.get("id"),
        "created_at": post.get("created_at"),
        "rating": post.get("rating"),
        "tags": {
            str(category): [str(tag)[:200] for tag in tags][:500]
            for category, tags in categories.items()
            if isinstance(tags, list)
        },
        "score": {
            "up": score.get("up"),
            "down": score.get("down"),
            "total": score.get("total"),
        },
        "fav_count": post.get("fav_count"),
        "comment_count": post.get("comment_count"),
        "sources": [
            str(source)[:500]
            for source in (post.get("sources") or [])
            if source
        ][:20],
        "description": str(post.get("description") or "")[:E621_DESCRIPTION_LIMIT],
        "pools": [
            int(pool)
            for pool in (post.get("pools") or [])
            if str(pool).isdigit()
        ][:50],
        "relationships": {
            "parent_id": relationships.get("parent_id"),
            "has_children": relationships.get("has_children"),
        },
        "file": {
            "md5": file_data.get("md5"),
            "ext": file_data.get("ext"),
            "size": file_data.get("size"),
            "width": file_data.get("width"),
            "height": file_data.get("height"),
            "url": file_data.get("url"),
        },
        "preview": {
            "url": (post.get("preview") or {}).get("url"),
        },
        "uploader_name": post.get("uploader_name"),
    }


def sanitize_iqdb_results(results):
    """Keep a compact, size-bounded copy of IQDB candidates."""
    cleaned = []
    for result in results[:10]:
        if not isinstance(result, dict):
            continue
        score = result.get("score")
        post_id = result.get("post_id")
        tags_preview = result.get("tags_preview")
        entry = {
            "post_id": post_id if isinstance(post_id, int) else None,
            "score": float(score) if isinstance(score, (int, float)) else None,
            "preview_url": str(result.get("preview_url") or "")[:500] or None,
            "rating": str(result.get("rating") or "")[:1] or None,
            "md5": str(result.get("md5") or "")[:32] or None,
            "score_total": (
                result.get("score_total")
                if isinstance(result.get("score_total"), int)
                else None
            ),
            "fav_count": (
                result.get("fav_count")
                if isinstance(result.get("fav_count"), int)
                else None
            ),
            "width": (
                result.get("width")
                if isinstance(result.get("width"), int)
                else None
            ),
            "height": (
                result.get("height")
                if isinstance(result.get("height"), int)
                else None
            ),
            "tags_preview": (
                [str(tag)[:100] for tag in tags_preview][:8]
                if isinstance(tags_preview, list)
                else []
            ),
        }
        if entry["post_id"] or entry["preview_url"]:
            cleaned.append(entry)
    return cleaned


def generate_video_thumbnail(md5, path):
    """Extract a JPEG thumbnail from a video, cached under MEDIA_ROOT/thumbs."""
    if not shutil.which("ffmpeg"):
        return None
    thumbs_dir = Path(settings.MEDIA_ROOT) / "thumbs"
    thumbs_dir.mkdir(parents=True, exist_ok=True)
    target = thumbs_dir / f"{md5}.jpg"
    if target.exists() and target.stat().st_mtime >= os.path.getmtime(path):
        return target
    command = [
        "ffmpeg",
        "-y",
        "-ss",
        "0.5",
        "-i",
        str(path),
        "-frames:v",
        "1",
        "-vf",
        "scale=480:-2",
        "-loglevel",
        "error",
        str(target),
    ]
    try:
        subprocess.run(command, check=True, capture_output=True, timeout=60)
    except (subprocess.SubprocessError, OSError):
        return None
    return target if target.exists() else None