New secrets system and quotas (#35)

extensions/exceptions.py

@@ -3,6 +3,7 @@ from fastapi.responses import UJSONResponse
 from starlette.status import (
     HTTP_400_BAD_REQUEST,
     HTTP_401_UNAUTHORIZED,
+    HTTP_403_FORBIDDEN,
     HTTP_404_NOT_FOUND,
     HTTP_406_NOT_ACCEPTABLE,
     HTTP_409_CONFLICT,
@@ -10,19 +11,20 @@ from starlette.status import (
 )
 
 from classes.exceptions import (
-    AlbumNotFoundError,
+    AccessTokenInvalidError,
     AlbumAlreadyExistsError,
     AlbumIncorrectError,
+    AlbumNotFoundError,
     PhotoNotFoundError,
     PhotoSearchQueryEmptyError,
-    VideoNotFoundError,
-    VideoSearchQueryEmptyError,
     SearchPageInvalidError,
     SearchTokenInvalidError,
-    AccessTokenInvalidError,
-    UserEmailCodeInvalid,
     UserAlreadyExists,
     UserCredentialsInvalid,
+    UserEmailCodeInvalid,
+    UserMediaQuotaReached,
+    VideoNotFoundError,
+    VideoSearchQueryEmptyError,
 )
 from modules.app import app
 
@@ -155,3 +157,13 @@ async def user_credentials_invalid_exception_handler(
         status_code=HTTP_401_UNAUTHORIZED,
         content={"detail": "Invalid credentials."},
     )
+
+
+@app.exception_handler(UserMediaQuotaReached)
+async def user_media_quota_reached_exception_handler(
+    request: Request, exc: UserMediaQuotaReached
+):
+    return UJSONResponse(
+        status_code=HTTP_403_FORBIDDEN,
+        content={"detail": "Media quota has been reached, media upload impossible."},
+    )

extensions/photos.py

@@ -30,6 +30,7 @@ from classes.exceptions import (
     SearchLimitInvalidError,
     SearchPageInvalidError,
     SearchTokenInvalidError,
+    UserMediaQuotaReached,
 )
 from classes.models import (
     Photo,
@@ -38,7 +39,7 @@ from classes.models import (
     SearchResultsPhoto,
 )
 from modules.app import app
-from modules.database import col_albums, col_photos, col_tokens
+from modules.database import col_albums, col_photos, col_tokens, col_videos
 from modules.exif_reader import extract_location
 from modules.hasher import get_duplicates, get_phash
 from modules.scheduler import scheduler
@@ -91,6 +92,7 @@ async def compress_image(image_path: str):
 
 
 photo_post_responses = {
+    403: UserMediaQuotaReached().openapi,
     404: AlbumNameNotFoundError("name").openapi,
     409: {
         "description": "Image Duplicates Found",
@@ -125,6 +127,13 @@ async def photo_upload(
     if (await col_albums.find_one({"user": current_user.user, "name": album})) is None:
         raise AlbumNameNotFoundError(album)
 
+    user_media_count = (
+        await col_photos.count_documents({"user": current_user.user})
+    ) + (await col_videos.count_documents({"user": current_user.user}))
+
+    if user_media_count >= current_user.quota and not current_user.quota == -1:  # type: ignore
+        raise UserMediaQuotaReached()
+
     makedirs(Path(f"data/users/{current_user.user}/albums/{album}"), exist_ok=True)
 
     filename = file.filename

extensions/users.py

@@ -109,6 +109,7 @@ if configGet("registration_enabled") is True:
             {
                 "user": user,
                 "email": email,
+                "quota": None,
                 "hash": get_password_hash(password),
                 "disabled": configGet("registration_requires_confirmation"),
             }

extensions/videos.py

@@ -21,6 +21,7 @@ from classes.exceptions import (
     SearchLimitInvalidError,
     SearchPageInvalidError,
     SearchTokenInvalidError,
+    UserMediaQuotaReached,
     VideoNotFoundError,
     VideoSearchQueryEmptyError,
 )
@@ -31,10 +32,13 @@ from classes.models import (
     VideoPublic,
 )
 from modules.app import app
-from modules.database import col_albums, col_tokens, col_videos
+from modules.database import col_albums, col_photos, col_tokens, col_videos
 from modules.security import User, get_current_active_user
 
-video_post_responses = {404: AlbumNameNotFoundError("name").openapi}
+video_post_responses = {
+    403: UserMediaQuotaReached().openapi,
+    404: AlbumNameNotFoundError("name").openapi,
+}
 
 
 @app.post(
@@ -53,6 +57,13 @@ async def video_upload(
     if (await col_albums.find_one({"user": current_user.user, "name": album})) is None:
         raise AlbumNameNotFoundError(album)
 
+    user_media_count = (
+        await col_videos.count_documents({"user": current_user.user})
+    ) + (await col_photos.count_documents({"user": current_user.user}))
+
+    if user_media_count >= current_user.quota and not current_user.quota == -1:  # type: ignore
+        raise UserMediaQuotaReached()
+
     makedirs(Path(f"data/users/{current_user.user}/albums/{album}"), exist_ok=True)
 
     filename = file.filename