New secrets system and quotas (#35)

.gitignore

@@ -153,5 +153,6 @@ cython_debug/
 #.idea/
 
 # Custom
-.vscode
+data/
+.vscode/
 config.json

README.md

@@ -47,7 +47,8 @@ First you need to have a Python interpreter, MongoDB and optionally git. You can
     1. Copy file `config_example.json` to `config.json`
     2. Open `config.json` using your favorite text editor. For example `nano config.json`
     3. Change `"database"` keys to match your MongoDB setup
-    4. Change `"external_address"` to the ip/http address you may get in responses. By default it's `"localhost"`. This is extremely useful when running behind reverse-proxy.  
+    4. Set the key `"secret"` to your JWT secret. You can type in anything, but long secrets are recommended. You can also set environment variable `PHOTOSAPI_SECRET` as an alternative
+    5. Change `"external_address"` to the ip/http address you may get in responses. By default it's `"localhost"`. This is extremely useful when running behind reverse-proxy.  
 
     After configuring everything listed above your API will be able to boot, however further configuration can be done. You can read about it in [repository's wiki](https://git.end-play.xyz/profitroll/PhotosAPI/wiki/Configuration). There's no need to focus on that now, it makes more sense to configure it afterwards.
 

classes/exceptions.py

@@ -286,3 +286,23 @@ class UserCredentialsInvalid(HTTPException):
             status_code=401,
             detail=self.openapi["content"]["application/json"]["example"]["detail"],
         )
+
+
+class UserMediaQuotaReached(HTTPException):
+    """Raises HTTP 403 if user's quota has been reached."""
+
+    def __init__(self):
+        self.openapi = {
+            "description": "Media Quota Reached",
+            "content": {
+                "application/json": {
+                    "example": {
+                        "detail": "Media quota has been reached, media upload impossible."
+                    }
+                }
+            },
+        }
+        super().__init__(
+            status_code=403,
+            detail=self.openapi["content"]["application/json"]["example"]["detail"],
+        )

config_example.json

@@ -6,6 +6,7 @@
         "user": null,
         "password": null
     },
+    "secret": "",
     "messages": {
         "email_confirmed": "Email confirmed. You can now log in."
     },
@@ -14,6 +15,7 @@
     "media_token_valid_hours": 12,
     "registration_enabled": true,
     "registration_requires_confirmation": false,
+    "default_user_quota": 10000,
     "mailer": {
         "smtp": {
             "host": "",

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

modules/app.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI
 from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
 
-app = FastAPI(title="END PLAY Photos", docs_url=None, redoc_url=None, version="0.5")
+app = FastAPI(title="END PLAY Photos", docs_url=None, redoc_url=None, version="0.6")
 
 
 @app.get("/docs", include_in_schema=False)

modules/security.py

@@ -1,4 +1,5 @@
 from datetime import datetime, timedelta, timezone
+from os import getenv
 from typing import List, Union
 
 from fastapi import Depends, HTTPException, Security, status
@@ -8,9 +9,26 @@ from passlib.context import CryptContext
 from pydantic import BaseModel, ValidationError
 
 from modules.database import col_users
+from modules.utils import configGet
+
+try:
+    configGet("secret")
+except KeyError as exc:
+    raise KeyError(
+        "PhotosAPI secret is not set. Secret key handling has changed in PhotosAPI 0.6.0, so you need to add the config key 'secret' to your config file."
+    ) from exc
+
+if configGet("secret") == "" and getenv("PHOTOSAPI_SECRET") is None:
+    raise KeyError(
+        "PhotosAPI secret is not set. Set the config key 'secret' or provide the environment variable 'PHOTOSAPI_SECRET' containing a secret string."
+    )
+
+SECRET_KEY = (
+    getenv("PHOTOSAPI_SECRET")
+    if getenv("PHOTOSAPI_SECRET") is not None
+    else configGet("secret")
+)
 
-with open("secret_key", "r", encoding="utf-8") as f:
-    SECRET_KEY = f.read()
 ALGORITHM = "HS256"
 ACCESS_TOKEN_EXPIRE_DAYS = 180
 
@@ -28,6 +46,7 @@ class TokenData(BaseModel):
 class User(BaseModel):
     user: str
     email: Union[str, None] = None
+    quota: Union[int, None] = None
     disabled: Union[bool, None] = None
 
 
@@ -71,6 +90,7 @@ async def get_user(user: str) -> UserInDB:
     return UserInDB(
         user=found_user["user"],
         email=found_user["email"],
+        quota=found_user["quota"],
         disabled=found_user["disabled"],
         hash=found_user["hash"],
     )
@@ -87,13 +107,16 @@ def create_access_token(
     data: dict, expires_delta: Union[timedelta, None] = None
 ) -> str:
     to_encode = data.copy()
+
     if expires_delta:
         expire = datetime.now(tz=timezone.utc) + expires_delta
     else:
         expire = datetime.now(tz=timezone.utc) + timedelta(
             days=ACCESS_TOKEN_EXPIRE_DAYS
         )
+
     to_encode["exp"] = expire
+
     return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
 
 
@@ -114,8 +137,10 @@ async def get_current_user(
     try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
         user: str = payload.get("sub")
+
         if user is None:
             raise credentials_exception
+
         token_scopes = payload.get("scopes", [])
         token_data = TokenData(scopes=token_scopes, user=user)
     except (JWTError, ValidationError) as exc:
@@ -133,6 +158,7 @@ async def get_current_user(
                 detail="Not enough permissions",
                 headers={"WWW-Authenticate": authenticate_value},
             )
+
     return user_record
 
 
@@ -141,4 +167,5 @@ async def get_current_active_user(
 ):
     if current_user.disabled:
         raise HTTPException(status_code=400, detail="Inactive user")
+
     return current_user