163 lines
5.3 KiB
Python
163 lines
5.3 KiB
Python
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
|