CostEvalPlatform/dependencies.py

# -*- coding: utf-8 -*-
# @Time    : 2024/11/19 下午7:33
# @Author  : 河瞬
# @FileName: dependencies.py
# @Software: PyCharm
from typing import Annotated
from fastapi import Depends, Request, HTTPException, Cookie, Response
from jose import jwt, JWTError
from database import engine
from sqlmodel import Session, select
from config import Settings
from models import User


def get_session():
    with Session(engine) as session:
        yield session


def get_settings():
    return Settings()


SessionDep = Annotated[Session, Depends(get_session)]

SettingsDep = get_settings()


def get_current_user(response: Response, session_token: Annotated[str | None, Cookie()] = None, db: SessionDep = None,
                     settings: SettingsDep = SettingsDep):
    if not session_token:
        response.set_cookie(key="session_token", value="", httponly=True)
        raise HTTPException(status_code=401, detail="Not authenticated", )

    try:
        payload = jwt.decode(session_token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
        user_id = payload.get("id")
        if user_id is None:
            response.set_cookie(key="session_token", value="", httponly=True)
            raise HTTPException(status_code=401, detail="Invalid token")
    except JWTError:
        response.set_cookie(key="session_token", value="", httponly=True)
        raise HTTPException(status_code=401, detail="Invalid token")

    user = db.exec(select(User).where(User.id == user_id)).first()
    if not user:
        response.set_cookie(key="session_token", value="", httponly=True)
        raise HTTPException(status_code=401, detail="User not found")

    return user