JakeBreath
27cbfd882c
- Active jobs on /stats get a cancel button wired to the existing download/match cancel endpoints, showing "cancelling..." and an inline error when the task already finished. - Cancelling now sets the status immediately, so a task whose runner died in a restart stops showing as "downloading". - Download streams use a bounded read timeout (10 s connect / 60 s read): a stalled socket fails within a minute (previously it could block forever), and a task cancelled while stalled is marked cancelled rather than error. - The stats job list reaps stale download/match tasks, so phantom jobs never appear on the dashboard.
576 lines
22 KiB
Python
576 lines
22 KiB
Python
import re
|
|
import shutil
|
|
import uuid
|
|
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, Q
|
|
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 e621, matching, services
|
|
from .downloads import reap_stale_downloads, start_download_task
|
|
from .matching import reap_stale_match_tasks, start_match_task
|
|
from .models import DownloadTask, MatchTask, MediaItem
|
|
from .permissions import CanUpload, IsUploaderOrStaffOrReadOnly
|
|
from .serializers import (
|
|
DownloadTaskSerializer,
|
|
MatchTaskSerializer,
|
|
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)
|
|
j_ids = self.request.query_params.get("j_ids", "").strip()
|
|
if j_ids:
|
|
numeric_ids = []
|
|
for value in j_ids.split(","):
|
|
text = value.strip()
|
|
number = text[2:] if text.upper().startswith("J-") else text
|
|
if number.isdigit():
|
|
numeric_ids.append(int(number))
|
|
queryset = queryset.filter(pk__in=numeric_ids)
|
|
search = self.request.query_params.get("search", "").strip()
|
|
search_type = self.request.query_params.get("search_type", "filename").strip()
|
|
if search:
|
|
words = [word for word in search.lower().split() if word]
|
|
if search_type == "tags":
|
|
for word in words:
|
|
queryset = queryset.filter(search_tags__icontains=word)
|
|
elif search_type == "both":
|
|
query = Q(locations__rel_path__icontains=search)
|
|
tag_query = Q()
|
|
for word in words:
|
|
tag_query &= Q(search_tags__icontains=word)
|
|
queryset = queryset.filter(query | tag_query)
|
|
else:
|
|
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)
|
|
statuses = []
|
|
for value in self.request.query_params.getlist("status"):
|
|
statuses.extend(
|
|
part.strip() for part in value.split(",") if part.strip()
|
|
)
|
|
if statuses:
|
|
status_query = Q()
|
|
for status_value in set(statuses):
|
|
if status_value == "custom":
|
|
status_query |= Q(has_custom_data=True)
|
|
elif status_value == "matched":
|
|
status_query |= Q(e621_match_status=MediaItem.E621_MATCHED)
|
|
elif status_value == "not_found":
|
|
status_query |= Q(e621_match_status=MediaItem.E621_NOT_FOUND)
|
|
elif status_value == "deleted":
|
|
status_query |= Q(e621_match_status=MediaItem.E621_DELETED)
|
|
elif status_value == "unknown":
|
|
status_query |= Q(
|
|
has_custom_data=False,
|
|
e621_match_status=MediaItem.E621_UNKNOWN,
|
|
)
|
|
queryset = queryset.filter(status_query)
|
|
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,
|
|
}
|
|
)
|
|
|
|
def _can_match(self, request, item):
|
|
user = request.user
|
|
return bool(
|
|
user.is_superuser
|
|
or user.role == user.ROLE_STAFF
|
|
or item.uploaded_by_id == user.id
|
|
)
|
|
|
|
@action(detail=True, methods=["post"], permission_classes=[CanUpload])
|
|
def match(self, request, pk=None):
|
|
"""Link this item to e621: MD5 lookup, or an explicit post id."""
|
|
item = self.get_object()
|
|
if not self._can_match(request, item):
|
|
return Response(
|
|
{"detail": "Only the uploader or staff can match this item."},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
extra = {}
|
|
post_id = request.data.get("post_id")
|
|
try:
|
|
if post_id not in (None, ""):
|
|
text = str(post_id).strip()
|
|
if not text.isdigit():
|
|
return Response(
|
|
{"detail": "Post ID must be a number."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
extra["md5_match"] = matching.link_item_to_post(
|
|
item, int(text), request.user
|
|
)
|
|
else:
|
|
matching.check_item(item, request.user)
|
|
except e621.E621NotFound:
|
|
return Response(
|
|
{"detail": "That post does not exist on e621."},
|
|
status=status.HTTP_404_NOT_FOUND,
|
|
)
|
|
except e621.E621Error as exc:
|
|
return Response(
|
|
{"detail": str(exc)}, status=status.HTTP_502_BAD_GATEWAY
|
|
)
|
|
item.refresh_from_db()
|
|
data = dict(self.get_serializer(item).data)
|
|
data.update(extra)
|
|
return Response(data)
|
|
|
|
@action(detail=True, methods=["post"], permission_classes=[CanUpload])
|
|
def unlink(self, request, pk=None):
|
|
"""Forget this item's e621 link so it can be matched again."""
|
|
item = self.get_object()
|
|
if not self._can_match(request, item):
|
|
return Response(
|
|
{"detail": "Only the uploader or staff can unlink this item."},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
matching.unlink_item(item)
|
|
return Response(self.get_serializer(item).data)
|
|
|
|
@action(detail=True, methods=["post"], permission_classes=[CanUpload])
|
|
def optimize(self, request, pk=None):
|
|
"""Replace this item's file(s) with a client-processed version.
|
|
|
|
The browser does the actual compression (WebCodecs/WASM); this only
|
|
applies the result. The item keeps its J-ID while MD5, size and
|
|
perceptual hashes are recomputed, so it may become a duplicate again.
|
|
"""
|
|
item = self.get_object()
|
|
if not self._can_match(request, item):
|
|
return Response(
|
|
{"detail": "Only the uploader or staff can optimize this item."},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
upload = request.FILES.get("file")
|
|
if upload is None:
|
|
return Response(
|
|
{"detail": "A processed file is required."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
extension = Path(upload.name).suffix.lower()
|
|
if extension not in services.ALLOWED_EXTENSIONS:
|
|
return Response(
|
|
{"detail": f"Unsupported file type: {extension or 'unknown'}"},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
locations = list(item.locations.all())
|
|
if not locations:
|
|
return Response(
|
|
{"detail": "This item has no file on disk."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
folder = Path(locations[0].path).parent
|
|
temp_path = folder / f".optimize-{uuid.uuid4().hex}{extension}"
|
|
try:
|
|
with open(temp_path, "wb") as handle:
|
|
for chunk in upload.chunks():
|
|
handle.write(chunk)
|
|
new_md5 = services.compute_md5(temp_path)
|
|
new_size = temp_path.stat().st_size
|
|
except OSError as exc:
|
|
temp_path.unlink(missing_ok=True)
|
|
return Response(
|
|
{"detail": f"Could not store the processed file: {exc}"},
|
|
status=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
|
|
if new_md5 == item.md5:
|
|
temp_path.unlink(missing_ok=True)
|
|
return Response(
|
|
{"detail": "The processed file is identical to the current one."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
clash = MediaItem.objects.filter(md5=new_md5).exclude(pk=item.pk).first()
|
|
if clash is not None:
|
|
temp_path.unlink(missing_ok=True)
|
|
return Response(
|
|
{
|
|
"detail": (
|
|
f"The processed file is an exact match (MD5) for "
|
|
f"J-{clash.id}; nothing was changed."
|
|
)
|
|
},
|
|
status=status.HTTP_409_CONFLICT,
|
|
)
|
|
|
|
new_name = f"J-{item.id}{extension}"
|
|
try:
|
|
for location in locations:
|
|
target = Path(location.path).parent / new_name
|
|
shutil.copyfile(temp_path, target)
|
|
if str(target) != location.path:
|
|
Path(location.path).unlink(missing_ok=True)
|
|
location.path = str(target)
|
|
try:
|
|
location.rel_path = str(
|
|
target.relative_to(settings.WATCHED_FOLDER)
|
|
)
|
|
except ValueError:
|
|
location.rel_path = target.name
|
|
location.mtime = target.stat().st_mtime
|
|
location.save(update_fields=["path", "rel_path", "mtime"])
|
|
except OSError as exc:
|
|
return Response(
|
|
{"detail": f"Could not replace the file(s): {exc}"},
|
|
status=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
finally:
|
|
temp_path.unlink(missing_ok=True)
|
|
|
|
hashes = (
|
|
services.compute_visual_hashes(locations[0].path)
|
|
if extension in services.IMAGE_EXTENSIONS
|
|
else {}
|
|
)
|
|
item.md5 = new_md5
|
|
item.size = new_size
|
|
for field in services.HASH_FIELDS:
|
|
setattr(item, field, hashes.get(field, ""))
|
|
item.save(
|
|
update_fields=["md5", "size", *services.HASH_FIELDS, "updated_at"]
|
|
)
|
|
|
|
item = self.get_queryset().get(pk=item.pk)
|
|
return Response(self.get_serializer(item).data)
|
|
|
|
|
|
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):
|
|
reap_stale_downloads()
|
|
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,
|
|
status=DownloadTask.STATUS_CANCELLED,
|
|
speed=None,
|
|
updated_at=timezone.now(),
|
|
)
|
|
return Response({"success": True})
|
|
|
|
|
|
class MatchTaskViewSet(
|
|
mixins.ListModelMixin,
|
|
mixins.RetrieveModelMixin,
|
|
viewsets.GenericViewSet,
|
|
):
|
|
"""Background library-wide e621 match scans."""
|
|
|
|
serializer_class = MatchTaskSerializer
|
|
permission_classes = [CanUpload]
|
|
http_method_names = ["get", "post", "head", "options"]
|
|
|
|
def get_queryset(self):
|
|
reap_stale_match_tasks()
|
|
queryset = MatchTask.objects.all()
|
|
user = self.request.user
|
|
if not (user.is_staff or user.is_superuser):
|
|
queryset = queryset.filter(user=user)
|
|
return queryset
|
|
|
|
def create(self, request):
|
|
scope = str(request.data.get("scope") or MatchTask.SCOPE_MISSING).strip()
|
|
if scope not in {MatchTask.SCOPE_MISSING, MatchTask.SCOPE_ALL}:
|
|
scope = MatchTask.SCOPE_MISSING
|
|
task = MatchTask.objects.create(user=request.user, scope=scope)
|
|
start_match_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 {MatchTask.STATUS_PENDING, MatchTask.STATUS_RUNNING}:
|
|
return Response(
|
|
{"detail": "Task is not running."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
MatchTask.objects.filter(id=task.id).update(
|
|
cancelled=True,
|
|
status=MatchTask.STATUS_CANCELLED,
|
|
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
|