from datetime import datetime, timedelta from passlib.context import CryptContext from fastapi import Depends, HTTPException, status from fastapi.security import OAuth2PasswordBearer from sqlalchemy.orm import joinedload from pydantic import BaseModel from jose import JWTError, jwt from sqlmodel import select from treetrail.config import conf from treetrail.models import User, Role, UserRoleLink from treetrail.database import db_session # openssl rand -hex 32 # import secrets # SECRET_KEY = secrets.token_hex(32) ALGORITHM: str = "HS256" class Token(BaseModel): access_token: str token_type: str class TokenData(BaseModel): username: str | None = None # class User(BaseModel): # username: str # email: str | None = None # full_name: str | None = None # disabled: bool = False oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token", auto_error=False) pwd_context = CryptContext(schemes=["sha256_crypt", "bcrypt"], deprecated="auto") credentials_exception = HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"}, ) def get_password_hash(password: str): return pwd_context.hash(password) async def delete_user(username) -> None: async with db_session() as session: user_in_db = await get_user(username) if user_in_db is None: raise SystemExit(f'User {username} does not exist in the database') await session.delete(user_in_db) async def enable_user(username, enable=True) -> None: async with db_session() as session: user_in_db = await get_user(username) if user_in_db is None: raise SystemExit(f'User {username} does not exist in the database') user_in_db.disabled = not enable # type: ignore session.add(user_in_db) await session.commit() async def create_user(username: str, password: str, full_name: str | None = None, email: str | None = None) -> User: async with db_session() as session: user = await get_user(username) if user is None: user = User( username=username, password=get_password_hash(password), full_name=full_name, email=email, disabled=False ) session.add(user) await session.commit() else: user.full_name = full_name # type: ignore user.email = email # type: ignore user.password = get_password_hash(password) # type: ignore await session.commit() await session.refresh(user) return user async def get_user(username: str) -> User | None: # type: ignore async with db_session() as session: query = select(User)\ .where(User.username==username)\ .options(joinedload(User.roles)) # type: ignore data = await session.exec(query) return data.first() def verify_password(plain_password, hashed_password): return pwd_context.verify(plain_password, hashed_password) async def get_current_user(token: str = Depends(oauth2_scheme)) -> User | None: # type: ignore if token is None or token == 'null': return None try: payload = jwt.decode(token, conf.security.secret_key, algorithms=[ALGORITHM]) username: str = payload.get("sub", '') if username == '': raise credentials_exception token_data = TokenData(username=username) except JWTError: return None user = await get_user(username=token_data.username) # type: ignore return user async def authenticate_user(username: str, password: str): user = await get_user(username) if not user: return False if not verify_password(password, user.password): return False return user async def get_current_active_user( current_user: User | None = Depends(get_current_user)) -> User: # type: ignore if current_user is None: raise HTTPException(status_code=400, detail="Not authenticated") if current_user.disabled: raise HTTPException(status_code=400, detail="Inactive user") return current_user async def get_current_roles(user: User | None = Depends(get_current_user)) -> list[Role]: # type: ignore roles: list[Role] if user is None: roles = [] else: roles = user.roles return roles def create_access_token(data: dict, expires_delta: timedelta): to_encode = data.copy() expire = datetime.utcnow() + expires_delta to_encode.update({"exp": expire}) encoded_jwt = jwt.encode(to_encode, conf.security.secret_key, algorithm=ALGORITHM) return encoded_jwt async def add_user_role(username: str, role_id: str): async with db_session() as session: user_in_db = await get_user(username) if user_in_db is None: raise SystemExit(f'User {username} does not exist in the database') user_role = UserRoleLink(user_id=user_in_db.username, role_id=role_id) session.add(user_role) await session.commit() async def add_role(role_id: str) -> Role: async with db_session() as session: role = Role(name=role_id) session.add(role) await session.commit() await session.refresh(role) return role