from datetime import timedelta, timezone, datetime from typing import Annotated from fastapi import Depends, HTTPException, Request, Response, status from pydantic import BaseModel, ValidationError import jwt from jwt.exceptions import ExpiredSignatureError, InvalidTokenError from sqlmodel import Session, select from db import engine, User from fastapi.security import ( OAuth2PasswordBearer, OAuth2PasswordRequestForm, SecurityScopes, ) from pydantic_settings import BaseSettings, SettingsConfigDict from passlib.context import CryptContext from sqlalchemy.exc import OperationalError class Config(BaseSettings): secret_key: str = "" access_token_expire_minutes: int = 15 model_config = SettingsConfigDict( env_file=".env", env_file_encoding="utf-8", extra="ignore" ) config = Config() class Token(BaseModel): access_token: str class TokenData(BaseModel): username: str | None = None scopes: list[str] = [] pwd_context = CryptContext(schemes=["argon2"], deprecated="auto") oauth2_scheme = OAuth2PasswordBearer( tokenUrl="api/token", scopes={ "analysis": "Access the results.", }, ) def verify_password(plain_password, hashed_password): return pwd_context.verify(plain_password, hashed_password) def get_password_hash(password): return pwd_context.hash(password) def get_user(username: str | None): if username: try: with Session(engine) as session: return session.exec( select(User).where(User.username == username) ).one_or_none() except OperationalError: return def authenticate_user(username: str, password: str): user = get_user(username) if not user: return False if not verify_password(password, user.hashed_password): return False return user def create_access_token(data: dict, expires_delta: timedelta | None = None): to_encode = data.copy() if expires_delta: expire = datetime.now(timezone.utc) + expires_delta else: expire = datetime.now(timezone.utc) + timedelta( seconds=config.access_token_expire_minutes ) to_encode.update({"exp": expire}) encoded_jwt = jwt.encode(to_encode, config.secret_key, algorithm="HS256") return encoded_jwt async def get_current_user(security_scopes: SecurityScopes, request: Request): if security_scopes.scopes: authenticate_value = f'Bearer scope="{security_scopes.scope_str}"' else: authenticate_value = "Bearer" credentials_exception = HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials", headers={"WWW-Authenticate": authenticate_value}, ) access_token = request.cookies.get("access_token") if not access_token: raise credentials_exception try: payload = jwt.decode(access_token, config.secret_key, algorithms=["HS256"]) username: str = payload.get("sub") if username is None: raise credentials_exception token_scopes = payload.get("scopes", []) token_data = TokenData(username=username, scopes=token_scopes) except ExpiredSignatureError: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Access token expired", headers={"WWW-Authenticate": authenticate_value}, ) except (InvalidTokenError, ValidationError): raise credentials_exception user = get_user(username=token_data.username) if user is None: raise credentials_exception for scope in security_scopes.scopes: if scope not in token_data.scopes: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Not enough permissions", headers={"WWW-Authenticate": authenticate_value}, ) return user async def get_current_active_user( current_user: Annotated[User, Depends(get_current_user)], ): if current_user.disabled: raise HTTPException(status_code=400, detail="Inactive user") return current_user async def login_for_access_token( form_data: Annotated[OAuth2PasswordRequestForm, Depends()], response: Response ) -> Token: user = authenticate_user(form_data.username, form_data.password) if not user: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect username or password", headers={"WWW-Authenticate": "Bearer"}, ) allowed_scopes = set(user.scopes.split()) requested_scopes = set(form_data.scopes) access_token = create_access_token( data={"sub": user.username, "scopes": list(allowed_scopes)} ) response.set_cookie( "access_token", value=access_token, httponly=True, samesite="strict", max_age=15, ) return Token(access_token=access_token) async def logout(response: Response): response.set_cookie("access_token", "", expires=0, httponly=True, samesite="strict") return {"message": "Successfully logged out"} async def read_users_me( current_user: Annotated[User, Depends(get_current_active_user)], ): return current_user async def read_own_items( current_user: Annotated[User, Depends(get_current_active_user)], ): return [{"item_id": "Foo", "owner": current_user.username}]