browsertrix/backend/btrixcloud/auth.py

""" auth functions for login """

import os
import uuid
from datetime import datetime, timedelta
from typing import Optional, Tuple, List
from passlib import pwd
from passlib.context import CryptContext

from pydantic import BaseModel
import jwt

from fastapi import (
    Request,
    HTTPException,
    Depends,
    WebSocket,
    APIRouter,
)

from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm

from .models import User


# ============================================================================
PASSWORD_SECRET = os.environ.get("PASSWORD_SECRET", uuid.uuid4().hex)

JWT_TOKEN_LIFETIME = int(os.environ.get("JWT_TOKEN_LIFETIME_MINUTES", 60)) * 60

ALGORITHM = "HS256"

RESET_VERIFY_TOKEN_LIFETIME_MINUTES = 60

PWD_CONTEXT = CryptContext(schemes=["bcrypt"], deprecated="auto")

# Audiences
AUTH_AUD = "btrix:auth"
RESET_AUD = "btrix:reset"
VERIFY_AUD = "btrix:verify"

# include fastapi-users audiences for backwards compatibility
AUTH_ALLOW_AUD = [AUTH_AUD, "fastapi-users:auth"]
RESET_ALLOW_AUD = [RESET_AUD, "fastapi-users:reset"]
VERIFY_ALLOW_AUD = [VERIFY_AUD, "fastapi-users:verify"]


# ============================================================================
class BearerResponse(BaseModel):
    """JWT Login Response"""

    access_token: str
    token_type: str


# ============================================================================
# pylint: disable=too-few-public-methods
class OA2BearerOrQuery(OAuth2PasswordBearer):
    """Override bearer check to also test query"""

    async def __call__(
        self, request: Request = None, websocket: WebSocket = None  # type: ignore
    ) -> str:
        param = None
        exc = None
        # use websocket as request if no request
        request = request or websocket  # type: ignore
        try:
            param = await super().__call__(request)  # type: ignore
            if param:
                return param

        # pylint: disable=broad-except
        except Exception as super_exc:
            exc = super_exc

        if request:
            param = request.query_params.get("auth_bearer")

        if param:
            return param

        if exc:
            raise exc

        raise HTTPException(status_code=404, detail="Not Found")


# ============================================================================
def generate_jwt(data: dict, minutes: int) -> str:
    """generate JWT token with expiration time (in minutes)"""
    expires_delta = timedelta(minutes=minutes)
    expire = datetime.utcnow() + expires_delta
    payload = data.copy()
    payload["exp"] = expire
    return jwt.encode(payload, PASSWORD_SECRET, algorithm=ALGORITHM)


# ============================================================================
def decode_jwt(token: str, audience: Optional[List[str]] = None) -> dict:
    """decode JWT token"""
    return jwt.decode(token, PASSWORD_SECRET, algorithms=[ALGORITHM], audience=audience)


# ============================================================================
def create_access_token(user: User) -> str:
    """get jwt token"""
    return generate_jwt({"sub": str(user.id), "aud": AUTH_AUD}, JWT_TOKEN_LIFETIME)


# ============================================================================
def verify_password(plain_password: str, hashed_password: str) -> bool:
    """verify password by hash"""
    return PWD_CONTEXT.verify(plain_password, hashed_password)


# ============================================================================
def verify_and_update_password(
    plain_password: str, hashed_password: str
) -> Tuple[bool, str]:
    """verify password and return updated hash, if any"""
    return PWD_CONTEXT.verify_and_update(plain_password, hashed_password)


# ============================================================================
def get_password_hash(password: str) -> str:
    """generate hash for password"""
    return PWD_CONTEXT.hash(password)


# ============================================================================
def generate_password() -> str:
    """generate new secure password"""
    return pwd.genword()


# ============================================================================
# pylint: disable=raise-missing-from
def init_jwt_auth(user_manager):
    """init jwt auth router + current_active_user dependency"""
    oauth2_scheme = OA2BearerOrQuery(tokenUrl="/api/auth/jwt/login", auto_error=False)

    async def get_current_user(token: str = Depends(oauth2_scheme)) -> User:
        credentials_exception = HTTPException(
            status_code=401,
            detail="invalid_credentials",
            headers={"WWW-Authenticate": "Bearer"},
        )
        try:
            payload = decode_jwt(token, AUTH_ALLOW_AUD)
            uid: Optional[str] = payload.get("sub") or payload.get("user_id")
            if uid is None:
                raise credentials_exception
        except Exception:
            raise credentials_exception
        user = await user_manager.get_by_id(uuid.UUID(uid))
        if user is None:
            raise credentials_exception
        return user

    current_active_user = get_current_user

    auth_jwt_router = APIRouter()

    def get_bearer_response(user: User):
        """get token, return bearer response for user"""
        token = create_access_token(user)
        return BearerResponse(access_token=token, token_type="bearer")

    @auth_jwt_router.post("/login", response_model=BearerResponse)
    async def login(
        credentials: OAuth2PasswordRequestForm = Depends(),
    ):
        user = await user_manager.authenticate(
            credentials.username, credentials.password
        )

        if not user:
            raise HTTPException(
                status_code=400,
                detail="login_bad_credentials",
            )
        # if requires_verification and not user.is_verified:
        #    raise HTTPException(
        #        status_code=400,
        #        detail="login_user_not_verified",
        #    )
        return get_bearer_response(user)

    @auth_jwt_router.post("/refresh", response_model=BearerResponse)
    async def refresh_jwt(user=Depends(current_active_user)):
        return get_bearer_response(user)

    return auth_jwt_router, current_active_user