kiwi-vpn/api/kiwi_vpn_api/routers/_common.py

"""
Common dependencies for routers.
"""


from fastapi import Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from sqlalchemy.orm import Session

from ..config import Config
from ..db import Connection
from ..db.schemas import User

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="user/authenticate")


class Responses:
    """
    Just a namespace.

    Describes API response status codes.
    """

    OK = {
        "content": None,
    }
    INSTALLED = {
        "description": "kiwi-vpn already installed",
        "content": None,
    }
    NOT_INSTALLED = {
        "description": "kiwi-vpn not installed",
        "content": None,
    }
    NEEDS_USER = {
        "description": "Must be logged in",
        "content": None,
    }
    NEEDS_ADMIN = {
        "description": "Must be admin",
        "content": None,
    }
    ENTRY_EXISTS = {
        "description": "Entry exists in database",
        "content": None,
    }


async def get_current_user(
    token: str = Depends(oauth2_scheme),
    db: Session | None = Depends(Connection.get),
    current_config: Config | None = Depends(Config.load),
) -> User | None:
    """
    Get the currently logged-in user from the database.
    """

    # can't connect to an unconfigured database
    if current_config is None:
        return None

    username = await current_config.jwt.decode_token(token)
    user = User.from_db(db, username)

    return user


async def get_current_user_if_admin(
    current_config: Config | None = Depends(Config.load),
    current_user: User | None = Depends(get_current_user),
) -> User:
    """
    Get the currently logged-in user if it is an admin.
    """

    # fail if not installed
    if current_config is None:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)

    # fail if not requested by a user
    if current_user is None:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)

    # fail if not requested by an admin
    if not current_user.is_admin():
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)


async def get_current_user_if_admin_or_self(
    user_name: str,
    current_config: Config | None = Depends(Config.load),
    current_user: User | None = Depends(get_current_user),
) -> User:
    """
    Get the currently logged-in user.

    Fails a) if the currently logged-in user is not the requested user,
    and b) if it is not an admin.
    """

    # fail if not installed
    if current_config is None:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)

    # fail if not requested by a user
    if current_user is None:
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)

    # fail if not requested by an admin or self
    if not (current_user.is_admin() or current_user.name == user_name):
        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)