classquiz-ai/classquiz/auth.py

#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
import uuid
from datetime import timedelta
from typing import Dict
from typing import Optional
from typing import Union

from fastapi import Depends
from fastapi import HTTPException
from fastapi import Request
from fastapi import status
from fastapi.openapi.models import OAuthFlows as OAuthFlowsModel
from fastapi.security import OAuth2
from fastapi.security.utils import get_authorization_scheme_param
from jose import JWTError, jwt
from passlib.hash import argon2

from classquiz.cache import get_cache
from classquiz.config import settings, redis
from datetime import datetime
from classquiz.db.models import User, TokenData, ApiKey

settings = settings()

pwd_context = argon2
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = settings.access_token_expire_minutes


class OAuth2PasswordBearerWithCookie(OAuth2):
    def __init__(
        self,
        tokenUrl: str,
        scheme_name: Optional[str] = None,
        scopes: Optional[Dict[str, str]] = None,
        auto_error: bool = True,
    ):
        if not scopes:
            scopes = {}
        flows = OAuthFlowsModel(password={"tokenUrl": tokenUrl, "scopes": scopes})
        super().__init__(flows=flows, scheme_name=scheme_name, auto_error=auto_error)

    async def __call__(self, request: Request) -> Optional[str]:
        try:
            authorization = request.state.access_token
        except AttributeError:
            authorization: str = request.cookies.get(
                "access_token"
            )  # changed to accept access token from httpOnly Cookie
        scheme, param = get_authorization_scheme_param(authorization)
        if not authorization or scheme.lower() != "bearer":
            if self.auto_error:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="Not authenticated",
                    headers={"WWW-Authenticate": "Bearer"},
                )
            else:
                return None
        return param


oauth2_scheme = OAuth2PasswordBearerWithCookie(tokenUrl="/api/v1/users/token/cookie")


def verify_password(plain_password: str, hashed_password: str) -> bool:
    return pwd_context.verify(plain_password, hashed_password)


def get_password_hash(password: str) -> str:
    return pwd_context.hash(password)


async def get_user_from_mail(email: str) -> Union[User, None]:
    return await get_cache(criteria="email", content=email)


async def get_user_from_username(username: str) -> Union[User, None]:
    return await get_cache(criteria="username", content=username)


async def get_user_from_id(id: str) -> Union[User, None]:
    return await get_cache(criteria="id", content=id)


async def authenticate_user(email: str, password: str) -> Union[User, bool]:
    user = await get_user_from_mail(email)
    if not user:
        return False
    if not verify_password(password, user.password):
        return False
    return user


def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, settings.secret_key, algorithm=ALGORITHM)
    return encoded_jwt


credentials_exception = HTTPException(
    status_code=status.HTTP_401_UNAUTHORIZED,
    detail="Could not validate credentials",
    headers={"WWW-Authenticate": "Bearer"},
)


async def get_current_user(token: str = Depends(oauth2_scheme)):
    try:
        payload = jwt.decode(token, settings.secret_key, algorithms=[ALGORITHM])
        email: str = payload.get("sub")
        if email is None:
            raise credentials_exception
        token_data = TokenData(email=email)
    except JWTError:
        raise credentials_exception
    user = await get_user_from_mail(email=token_data.email)
    if user is None:
        raise credentials_exception
    return user


async def get_admin_user(token: str = Depends(oauth2_scheme)) -> User:
    user = await get_current_user(token)
    admin_user = await User.objects.order_by(User.created_at.asc()).get()
    if admin_user.id == user.id:
        return user
    else:
        raise credentials_exception


async def get_current_user_optional(token: str = Depends(oauth2_scheme)) -> User | None:
    try:
        payload = jwt.decode(token, settings.secret_key, algorithms=[ALGORITHM])
        email: str = payload.get("sub")
        if email is None:
            return None
        token_data = TokenData(email=email)
    except JWTError:
        return None
    user = await get_user_from_mail(email=token_data.email)
    if user is None:
        return None
    return user


async def check_token(token: str = Depends(oauth2_scheme)):
    try:
        payload = jwt.decode(token, settings.secret_key, algorithms=[ALGORITHM])
        email: str = payload.get("sub")
        if email is None:
            raise credentials_exception
        token_data = TokenData(email=email)
    except JWTError:
        raise credentials_exception
    return token_data.email


async def check_api_key(key: str) -> uuid.UUID | None:
    redis_res = await redis.get(f"apikey:{key}")
    if redis_res is None:
        key2 = await ApiKey.objects.get_or_none(key=key)
        if key2 is None:
            return None
        else:
            await redis.set(f"apikey:{key}", key2.user.id.hex, ex=3600)
            return key2.user.id
    else:
        return uuid.UUID(redis_res)