Files
thePlug/app/core/security.py
T

110 lines
2.9 KiB
Python

from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import Annotated
from uuid import uuid4
from authlib.jose import jwt # type: ignore
from authlib.jose.errors import JoseError # type: ignore
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pwdlib import PasswordHash
from pwdlib.hashers.argon2 import Argon2Hasher
from core.config import settings
from models.auth import AccessTokenPayload, AuthUserDocument
from repositories.auth import get_user_by_username, is_token_revoked
pwd_context = PasswordHash((Argon2Hasher(),))
bearer_scheme = HTTPBearer(auto_error=False)
@dataclass(slots=True)
class AuthContext:
token: str
user: AuthUserDocument
payload: AccessTokenPayload
class InvalidTokenError(Exception):
pass
def token_auth_error() -> HTTPException:
return HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Could not validate credentials",
headers={"WWW-Authenticate": "Bearer"},
)
def hash_password(password: str) -> str:
return pwd_context.hash(password)
def verify_password(plain_password: str, password_hash: str) -> bool:
return pwd_context.check(plain_password, password_hash)
def create_access_token(*, subject: str) -> str:
now = datetime.now(UTC)
expires_delta = timedelta(minutes=settings.JWT_EXPIRE_MINUTES)
expires_at = int((now + expires_delta).timestamp())
header = {
"alg": settings.JWT_ALGORITHM,
"typ": "JWT",
}
payload = {
"sub": subject,
"jti": str(uuid4()),
"exp": expires_at,
}
token = jwt.encode(header, payload, settings.JWT_SECRET_KEY)
if isinstance(token, bytes):
return token.decode("utf-8")
return token
def decode_access_token(token: str) -> AccessTokenPayload:
try:
claims = jwt.decode(token, settings.JWT_SECRET_KEY)
claims.validate()
return AccessTokenPayload.model_validate(dict(claims))
except (JoseError, ValueError, TypeError):
raise InvalidTokenError from None
async def get_auth_context(
credentials: Annotated[
HTTPAuthorizationCredentials | None,
Depends(bearer_scheme),
],
) -> AuthContext:
if credentials is None or credentials.scheme.lower() != "bearer":
raise token_auth_error()
token = credentials.credentials
try:
payload = decode_access_token(token)
except InvalidTokenError:
raise token_auth_error() from None
if await is_token_revoked(payload.jti):
raise token_auth_error()
user = await get_user_by_username(payload.sub)
if user is None:
raise token_auth_error()
return AuthContext(token=token, user=user, payload=payload)
async def require_authenticated_user(
auth_context: Annotated[AuthContext, Depends(get_auth_context)],
) -> AuthUserDocument:
return auth_context.user