From f1a190f03018b371f302230c1dbaeec92ce9f71a Mon Sep 17 00:00:00 2001 From: Profitroll <47523801+profitrollgame@users.noreply.github.com> Date: Sat, 18 Feb 2023 00:19:46 +0100 Subject: [PATCH] Added access tokens for duplicates --- classes/exceptions.py | 14 ++++++++++ extensions/exceptions.py | 7 +++++ extensions/photos.py | 57 ++++++++++++++++++++++++++++++++++++---- 3 files changed, 73 insertions(+), 5 deletions(-) diff --git a/classes/exceptions.py b/classes/exceptions.py index 926cf8a..cea98d6 100644 --- a/classes/exceptions.py +++ b/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): diff --git a/extensions/exceptions.py b/extensions/exceptions.py index 793c2f6..df8caae 100644 --- a/extensions/exceptions.py +++ b/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( diff --git a/extensions/photos.py b/extensions/photos.py index f30e8f1..bb62c58 100644 --- a/extensions/photos.py +++ b/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 }