Added access tokens for duplicates

classes/exceptions.py

@@ -176,6 +176,20 @@ class UserAlreadyExists(Exception):
             }
         }
 
+class AccessTokenInvalidError(Exception):
+    """Raises HTTP 401 if access token is not valid."""
+    def __init__(self):
+        self.openapi = {
+            "description": "Invalid Access Token",
+            "content": {
+                "application/json": {
+                    "example": {
+                        "detail": "Invalid access token."
+                    }
+                }
+            }
+        }
+
 class UserCredentialsInvalid(Exception):
     """Raises HTTP 401 if user credentials are not valid."""
     def __init__(self):

extensions/exceptions.py

@@ -81,6 +81,13 @@ async def user_already_exists_exception_handler(request: Request, exc: UserAlrea
         content={"detail": "User with this username already exists."},
     )
 
+@app.exception_handler(AccessTokenInvalidError)
+async def access_token_invalid_exception_handler(request: Request, exc: AccessTokenInvalidError):
+    return UJSONResponse(
+        status_code=HTTP_401_UNAUTHORIZED,
+        content={"detail": "Invalid access token."},
+    )
+
 @app.exception_handler(UserCredentialsInvalid)
 async def user_credentials_invalid_exception_handler(request: Request, exc: UserCredentialsInvalid):
     return UJSONResponse(

extensions/photos.py

@@ -7,22 +7,26 @@ from typing import Union
 from magic import Magic
 from datetime import datetime, timedelta, timezone
 from os import makedirs, path, remove, system
-from classes.exceptions import AlbumNameNotFoundError, PhotoNotFoundError, PhotoSearchQueryEmptyError, SearchPageInvalidError, SearchTokenInvalidError
+
+from pydantic import ValidationError
+from classes.exceptions import AccessTokenInvalidError, AlbumNameNotFoundError, PhotoNotFoundError, PhotoSearchQueryEmptyError, SearchPageInvalidError, SearchTokenInvalidError
 from classes.models import Photo, PhotoPublic, SearchResultsPhoto
 from modules.exif_reader import extract_location
 from modules.hasher import get_phash, get_duplicates
 from modules.scheduler import scheduler
-from modules.security import User, get_current_active_user
+from modules.security import ALGORITHM, SECRET_KEY, TokenData, User, create_access_token, get_current_active_user, get_user
 from modules.app import app
 from modules.database import col_photos, col_albums, col_tokens
 from pymongo import DESCENDING
 from bson.objectid import ObjectId
 from bson.errors import InvalidId
 from plum.exceptions import UnpackError
+from jose import JWTError, jwt
 
 from fastapi import UploadFile, Security
 from fastapi.responses import UJSONResponse, Response
-from starlette.status import HTTP_204_NO_CONTENT, HTTP_409_CONFLICT
+from fastapi.exceptions import HTTPException
+from starlette.status import HTTP_204_NO_CONTENT, HTTP_401_UNAUTHORIZED, HTTP_409_CONFLICT
 
 from modules.utils import logWrite
 
@@ -60,7 +64,8 @@ photo_post_responses = {
                     "detail": "Image duplicates found. Pass 'ignore_duplicates=true' to ignore.",
                     "duplicates": [
                         "string"
-                    ]
+                    ],
+                    "access_token": "string"
                 }
             }
         }
@@ -88,10 +93,14 @@ async def photo_upload(file: UploadFile, album: str, ignore_duplicates: bool = F
     duplicates = await get_duplicates(file_hash, album)
 
     if len(duplicates) > 0 and ignore_duplicates is False:
+        duplicates_ids = []
+        for entry in duplicates:
+            duplicates_ids.append(entry["id"])
         return UJSONResponse(
             {
                 "detail": "Image duplicates found. Pass 'ignore_duplicates=true' to ignore.",
-                "duplicates": duplicates
+                "duplicates": duplicates,
+                "access_token": create_access_token(data={"sub": current_user.user, "scopes": ["me", "photos.read"], "allowed": duplicates_ids}, expires_delta=timedelta(hours=1))
             },
             status_code=HTTP_409_CONFLICT
         )
@@ -136,6 +145,44 @@ async def photo_upload(file: UploadFile, album: str, ignore_duplicates: bool = F
         }
     )
 
+photo_get_token_responses = {
+    401: AccessTokenInvalidError().openapi,
+    404: PhotoNotFoundError("id").openapi
+}
+@app.get("/photos/{id}/token/{token}", description="Get a photo by id", responses=photo_get_token_responses)
+async def photo_get_token(id: str, token: str):
+
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        user: str = payload.get("sub")
+        if user is None:
+            raise AccessTokenInvalidError()
+        token_scopes = payload.get("scopes", [])
+        token_data = TokenData(scopes=token_scopes, user=user)
+    except (JWTError, ValidationError) as exp:
+        print(exp, flush=True)
+        raise AccessTokenInvalidError()
+        
+    user = get_user(user=token_data.user)
+
+    if id not in payload.get("allowed", []):
+        raise AccessTokenInvalidError()
+
+    try:
+        image = col_photos.find_one( {"_id": ObjectId(id)} )
+        if image is None:
+            raise InvalidId(id)
+    except InvalidId:
+        raise PhotoNotFoundError(id)
+
+    image_path = path.join("data", "users", user.user, "albums", image["album"], image["filename"])
+
+    mime = Magic(mime=True).from_file(image_path)
+
+    with open(image_path, "rb") as f: image_file = f.read() 
+
+    return Response(image_file, media_type=mime)
+
 photo_get_responses = {
     404: PhotoNotFoundError("id").openapi
 }