toolshed/backend/files/media_urls.py

import io
import os
from datetime import timedelta

from django.core.files.base import ContentFile
from django.core.files.storage import default_storage
from django.http import HttpResponse
from django.urls import path
from django.db.models import Q
from django.conf import settings
from django.utils.http import http_date
from django.utils.timezone import now
from drf_yasg.utils import swagger_auto_schema
from PIL import Image
from rest_framework import status
from rest_framework.decorators import api_view, permission_classes, authentication_classes
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response

from authentication.signature_auth import SignatureAuthentication
from files.models import File

THUMBNAIL_SIZES = (32, 64, 256)


def _accessible_files(request):
    # Shared by media_urls and thumbnail_urls: a file is visible if the requester is
    # friends-or-self with whatever references it (item, profile picture), or it's their own staged photo.
    return File.objects.filter(
        Q(connected_items__owner__in=request.user.friends_or_self()) |
        Q(profile_picture_users__in=request.user.friends_or_self()) |
        Q(staged_by_workflows__owner__in=request.user.user.all())
    ).distinct()


def _cache_headers(etag):
    # Content is hash-addressed and can never change under a given URL, so it's cacheable forever.
    return {
        'ETag': etag,
        'Cache-Control': 'max-age=31536000, private, immutable',
        'Expires': http_date((now() + timedelta(days=365)).timestamp()),
    }


@swagger_auto_schema(method='GET', auto_schema=None)
@api_view(['GET'])
@permission_classes([IsAuthenticated])
@authentication_classes([SignatureAuthentication])
def media_urls(request, hash_path):
    # CORS is added automatically by middleware, except via X-Accel-Redirect, where nginx's
    # /redirect_media/ block must set it instead.
    #
    # Looked up by the derived storage path, not the raw hash, to match FileSerializer.name
    # (used for AuthenticatedImage's `src`) and the existing test suite (MediaUrlTestCase).
    try:
        file = _accessible_files(request).get(file=hash_path)

        # Must run before this check, else a bare hash + If-None-Match would let anyone probe
        # file existence for files they can't see.
        if request.META.get('HTTP_IF_NONE_MATCH') == file.hash:
            return HttpResponse(status=status.HTTP_304_NOT_MODIFIED)

        cache_headers = _cache_headers(file.hash)

        if settings.SERVE_X_ACCEL_REDIRECT:
            return HttpResponse(status=status.HTTP_200_OK,
                                content_type=file.mime_type,
                                headers={
                                    'X-Accel-Redirect': f'/redirect_media/{hash_path}',
                                    **cache_headers,
                                })
        else:
            # Reads via FieldFile.open() (not file.file.path) since tests swap in an in-memory storage backend.
            with file.file.open('rb') as fh:
                content = fh.read()
            return HttpResponse(status=status.HTTP_200_OK,
                                content_type=file.mime_type,
                                headers=cache_headers,
                                content=content)

    except File.DoesNotExist:
        return Response(status=status.HTTP_404_NOT_FOUND)


def _thumbnail_rel_path(file_hash, size):
    # Mirrors hash_upload()'s sharding under thumbnails/<size>/, reachable via the same nginx alias as originals.
    return os.path.join('thumbnails', str(size), file_hash[:2], file_hash[2:4], file_hash[4:6],
                         file_hash[6:] + '.jpg')


@swagger_auto_schema(method='GET', auto_schema=None)
@api_view(['GET'])
@permission_classes([IsAuthenticated])
@authentication_classes([SignatureAuthentication])
def thumbnail_urls(request, size, hash_path):
    if size not in THUMBNAIL_SIZES:
        return Response(status=status.HTTP_404_NOT_FOUND)

    try:
        file = _accessible_files(request).get(file=hash_path)

        etag = f'{file.hash}_{size}'
        if request.META.get('HTTP_IF_NONE_MATCH') == etag:
            return HttpResponse(status=status.HTTP_304_NOT_MODIFIED)

        # Read/write via default_storage, not a hand-rolled path, to work with both real-disk
        # and in-memory test storage.
        rel_path = _thumbnail_rel_path(file.hash, size)
        if not default_storage.exists(rel_path):
            # Always re-encoded as JPEG regardless of original format - simpler than preserving transparency at this scale.
            with file.file.open('rb') as fh:
                image = Image.open(fh)
                image.thumbnail((size, size))
                # Flatten through RGBA before dropping to RGB. See docs/implementation.md#rgba-flattening-avoids-revealing-black-under-transparent-pixels.
                rgba = image.convert('RGBA')
                flattened = Image.new('RGB', rgba.size, (255, 255, 255))
                flattened.paste(rgba, mask=rgba.getchannel('A'))
                buffer = io.BytesIO()
                flattened.save(buffer, 'JPEG', quality=90)
            default_storage.save(rel_path, ContentFile(buffer.getvalue()))

        cache_headers = _cache_headers(etag)

        if settings.SERVE_X_ACCEL_REDIRECT:
            return HttpResponse(status=status.HTTP_200_OK,
                                content_type='image/jpeg',
                                headers={
                                    'X-Accel-Redirect': f'/redirect_media/{rel_path}',
                                    **cache_headers,
                                })
        else:
            with default_storage.open(rel_path, 'rb') as fh:
                content = fh.read()
            return HttpResponse(status=status.HTTP_200_OK,
                                content_type='image/jpeg',
                                headers=cache_headers,
                                content=content)

    except File.DoesNotExist:
        return Response(status=status.HTTP_404_NOT_FOUND)


urlpatterns = [
    path('<int:size>/<path:hash_path>/', thumbnail_urls),
    path('<path:hash_path>', media_urls),
]