J621/backend/apps/library/views.py

import re
from pathlib import Path
from urllib.parse import urlparse

from django.conf import settings
from django.core import signing
from django.db.models import Min
from django.http import Http404, StreamingHttpResponse
from django.shortcuts import get_object_or_404
from django.utils import timezone
from django.utils.text import get_valid_filename
from rest_framework import mixins, status, viewsets
from rest_framework.decorators import action
from rest_framework.permissions import AllowAny, IsAuthenticatedOrReadOnly
from rest_framework.response import Response
from rest_framework.views import APIView

from . import services
from .downloads import start_download_task
from .models import DownloadTask, MediaItem
from .permissions import CanUpload, IsUploaderOrStaffOrReadOnly
from .serializers import DownloadTaskSerializer, MediaItemSerializer

LIST_ORDERINGS = {"name", "-name", "size", "-size", "created_at", "-created_at"}
MD5_RE = re.compile(r"[0-9a-fA-F]{32}")


class MediaItemViewSet(
    mixins.ListModelMixin,
    mixins.RetrieveModelMixin,
    mixins.UpdateModelMixin,
    viewsets.GenericViewSet,
):
    serializer_class = MediaItemSerializer
    permission_classes = [IsAuthenticatedOrReadOnly, IsUploaderOrStaffOrReadOnly]
    http_method_names = ["get", "post", "patch", "head", "options"]

    def get_queryset(self):
        queryset = MediaItem.objects.prefetch_related("locations").annotate(
            name=Min("locations__rel_path")
        )
        if not self.request.user.is_authenticated:
            queryset = queryset.filter(hidden_from_guests=False)
        search = self.request.query_params.get("search", "").strip()
        if search:
            queryset = queryset.filter(locations__rel_path__icontains=search)
        ratings = [
            value
            for value in self.request.query_params.get("rating", "").split(",")
            if value in {"s", "q", "e"}
        ]
        if ratings:
            queryset = queryset.filter(rating__in=ratings)
        ordering = self.request.query_params.get("ordering", "").strip()
        queryset = queryset.order_by(
            ordering if ordering in LIST_ORDERINGS else "-created_at"
        )
        return queryset.distinct()

    def get_object(self):
        """Resolve J-<id>, a bare numeric id, or an MD5 fingerprint."""
        value = str(self.kwargs.get("pk", "")).strip()
        queryset = self.get_queryset()
        if MD5_RE.fullmatch(value):
            obj = get_object_or_404(queryset, md5=value.lower())
        else:
            numeric = value[2:] if value.upper().startswith("J-") else value
            if not numeric.isdigit():
                raise Http404
            obj = get_object_or_404(queryset, pk=int(numeric))
        self.check_object_permissions(self.request, obj)
        return obj

    def _signed_media_item(self, request, action_name):
        """Signed media URLs let <img>/<video> tags bypass guest filtering."""
        if request.user.is_authenticated:
            return None
        signature = request.query_params.get("sig")
        if not signature:
            return None
        try:
            payload = signing.loads(
                signature, salt=services.MEDIA_FILE_SALT, max_age=86400
            )
        except signing.BadSignature:
            return None
        if payload.get("action") != action_name:
            return None
        return MediaItem.objects.filter(pk=payload.get("item")).first()

    def _media_object(self, request, action_name):
        item = self._signed_media_item(request, action_name)
        if item is not None:
            return item
        return self.get_object()

    @action(detail=True, methods=["get"])
    def raw(self, request, pk=None):
        item = self._media_object(request, "raw")
        location = item.locations.first()
        if location is None:
            return Response(
                {"detail": "No file for this item."},
                status=status.HTTP_404_NOT_FOUND,
            )
        return services.serve_file(
            request, location.path, download=request.query_params.get("download") == "1"
        )

    @action(detail=True, methods=["get"])
    def thumbnail(self, request, pk=None):
        item = self._media_object(request, "thumbnail")
        location = item.locations.first()
        if location is None:
            return Response(
                {"detail": "No file for this item."},
                status=status.HTTP_404_NOT_FOUND,
            )
        path = Path(location.path)
        if path.suffix.lower() in services.VIDEO_EXTENSIONS:
            thumbnail = services.generate_video_thumbnail(item.md5, path)
            if thumbnail is None:
                return Response(
                    {"detail": "Thumbnail unavailable."},
                    status=status.HTTP_404_NOT_FOUND,
                )
            return services.serve_file(request, thumbnail)
        return services.serve_file(request, path)

    @action(detail=False, methods=["post"], permission_classes=[AllowAny])
    def lookup(self, request):
        """Report which of the given MD5s exist in the library."""
        md5s = request.data.get("md5s")
        if not isinstance(md5s, list):
            return Response(
                {"detail": "md5s must be a list."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        wanted = {str(value).strip().lower()[:32] for value in md5s if value}
        queryset = MediaItem.objects.filter(md5__in=wanted)
        if not request.user.is_authenticated:
            queryset = queryset.filter(hidden_from_guests=False)
        rows = list(queryset.values_list("md5", "id"))
        return Response(
            {
                "found": sorted(md5 for md5, _ in rows),
                "j_ids": {md5: f"J-{item_id}" for md5, item_id in rows},
            }
        )

    @action(detail=True, methods=["get"])
    def neighbors(self, request, pk=None):
        """Previous/next items in the ordered list, for keyboard navigation."""
        item = self.get_object()
        ordering = request.query_params.get("ordering", "-created_at").strip()
        if ordering not in LIST_ORDERINGS:
            ordering = "-created_at"
        field = ordering.lstrip("-")
        descending = ordering.startswith("-")
        value = {
            "size": item.size,
            "created_at": item.created_at,
            "name": item.name,
        }[field]

        queryset = self.get_queryset().exclude(pk=item.pk)
        if descending:
            previous_query = queryset.filter(**{f"{field}__gt": value}).order_by(field)
            next_query = queryset.filter(**{f"{field}__lt": value}).order_by(f"-{field}")
        else:
            previous_query = queryset.filter(**{f"{field}__lt": value}).order_by(
                f"-{field}"
            )
            next_query = queryset.filter(**{f"{field}__gt": value}).order_by(field)

        def brief(obj):
            location = obj.locations.first()
            return {
                "j_id": f"J-{obj.id}",
                "md5": obj.md5,
                "filename": Path(location.rel_path).name if location else obj.md5,
            }

        previous = previous_query.first()
        next_item = next_query.first()
        return Response(
            {
                "previous": brief(previous) if previous else None,
                "next": brief(next_item) if next_item else None,
            }
        )


class DownloadTaskViewSet(
    mixins.ListModelMixin,
    mixins.RetrieveModelMixin,
    viewsets.GenericViewSet,
):
    """Background 'Download to Library' jobs."""

    serializer_class = DownloadTaskSerializer
    permission_classes = [CanUpload]
    http_method_names = ["get", "post", "head", "options"]

    def get_queryset(self):
        queryset = DownloadTask.objects.select_related("library_item")
        user = self.request.user
        if not (user.is_staff or user.is_superuser):
            queryset = queryset.filter(user=user)
        return queryset

    def create(self, request):
        url = str(request.data.get("url") or "").strip()
        post_id = request.data.get("post_id")
        filename = str(request.data.get("filename") or "").strip()
        if not url.startswith(("http://", "https://")):
            return Response(
                {"detail": "A valid file URL is required."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        trimmed = services.trim_e621_post(request.data.get("post"))
        resolved_post_id = None
        if trimmed is not None and trimmed.get("id") is not None:
            resolved_post_id = int(trimmed["id"])
        elif post_id is not None and str(post_id).isdigit():
            resolved_post_id = int(post_id)

        task = DownloadTask.objects.create(
            user=request.user,
            post_id=resolved_post_id,
            url=url,
            filename=filename,
            e621_data=trimmed,
        )
        start_download_task(task.id)
        return Response(
            self.get_serializer(task).data, status=status.HTTP_202_ACCEPTED
        )

    @action(detail=True, methods=["post"])
    def cancel(self, request, pk=None):
        task = self.get_object()
        if task.status not in {
            DownloadTask.STATUS_PENDING,
            DownloadTask.STATUS_DOWNLOADING,
        }:
            return Response(
                {"detail": "Task is not running."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        DownloadTask.objects.filter(id=task.id).update(
            cancelled=True, updated_at=timezone.now()
        )
        return Response({"success": True})


class ClientDownloadView(APIView):
    """Stream an e621 file straight to the browser (no library write)."""

    permission_classes = [AllowAny]

    def get(self, request):
        url = str(request.query_params.get("url") or "").strip()
        filename = str(request.query_params.get("filename") or "").strip()
        parsed = urlparse(url)
        if (
            parsed.scheme not in {"http", "https"}
            or parsed.hostname not in settings.E621_MEDIA_HOSTS
        ):
            return Response(
                {"detail": "URL not allowed."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        import requests

        try:
            upstream = requests.get(
                url,
                headers={"User-Agent": settings.USER_AGENT},
                stream=True,
                timeout=60,
            )
            upstream.raise_for_status()
        except requests.RequestException as exc:
            return Response(
                {"detail": f"Could not fetch the file: {exc}"},
                status=status.HTTP_502_BAD_GATEWAY,
            )

        content_type = upstream.headers.get(
            "Content-Type", "application/octet-stream"
        )
        name = get_valid_filename(
            filename or Path(parsed.path).name or "download"
        )

        def stream():
            try:
                for chunk in upstream.iter_content(chunk_size=64 * 1024):
                    if chunk:
                        yield chunk
            finally:
                upstream.close()

        response = StreamingHttpResponse(stream(), content_type=content_type)
        response["Content-Disposition"] = f'attachment; filename="{name}"'
        length = upstream.headers.get("Content-Length")
        if length:
            response["Content-Length"] = length
        return response