auths.py

from pydantic import BaseModel
from typing import Optional
import uuid
import logging
from sqlalchemy import String, Column, Boolean
from sqlalchemy.orm import Session

from apps.webui.models.users import UserModel, Users
from utils.utils import verify_password

from apps.webui.internal.db import Base

from config import SRC_LOG_LEVELS

log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])

####################
# DB MODEL
####################


class Auth(Base):
    __tablename__ = "auth"

    id = Column(String, primary_key=True)
    email = Column(String)
    password = Column(String)
    active = Column(Boolean)


class AuthModel(BaseModel):
    id: str
    email: str
    password: str
    active: bool = True


####################
# Forms
####################


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


class ApiKey(BaseModel):
    api_key: Optional[str] = None


class UserResponse(BaseModel):
    id: str
    email: str
    name: str
    role: str
    profile_image_url: str


class SigninResponse(Token, UserResponse):
    pass


class SigninForm(BaseModel):
    email: str
    password: str


class ProfileImageUrlForm(BaseModel):
    profile_image_url: str


class UpdateProfileForm(BaseModel):
    profile_image_url: str
    name: str


class UpdatePasswordForm(BaseModel):
    password: str
    new_password: str


class SignupForm(BaseModel):
    name: str
    email: str
    password: str
    profile_image_url: Optional[str] = "/user.png"


class AddUserForm(SignupForm):
    role: Optional[str] = "pending"


class AuthsTable:

    def insert_new_auth(
        self,
        db: Session,
        email: str,
        password: str,
        name: str,
        profile_image_url: str = "/user.png",
        role: str = "pending",
        oauth_sub: Optional[str] = None,
    ) -> Optional[UserModel]:
        log.info("insert_new_auth")

        id = str(uuid.uuid4())

        auth = AuthModel(
            **{"id": id, "email": email, "password": password, "active": True}
        )
        result = Auth(**auth.model_dump())
        db.add(result)

        user = Users.insert_new_user(
            db, id, name, email, profile_image_url, role, oauth_sub
        )

        db.commit()
        db.refresh(result)

        if result and user:
            return user
        else:
            return None

    def authenticate_user(
        self, db: Session, email: str, password: str
    ) -> Optional[UserModel]:
        log.info(f"authenticate_user: {email}")
        try:
            auth = db.query(Auth).filter_by(email=email, active=True).first()
            if auth:
                if verify_password(password, auth.password):
                    user = Users.get_user_by_id(db, auth.id)
                    return user
                else:
                    return None
            else:
                return None
        except:
            return None

    def authenticate_user_by_api_key(
        self, db: Session, api_key: str
    ) -> Optional[UserModel]:
        log.info(f"authenticate_user_by_api_key: {api_key}")
        # if no api_key, return None
        if not api_key:
            return None

        try:
            user = Users.get_user_by_api_key(db, api_key)
            return user if user else None
        except:
            return False

    def authenticate_user_by_trusted_header(
        self, db: Session, email: str
    ) -> Optional[UserModel]:
        log.info(f"authenticate_user_by_trusted_header: {email}")
        try:
            auth = db.query(Auth).filter(email=email, active=True).first()
            if auth:
                user = Users.get_user_by_id(auth.id)
                return user
        except:
            return None

    def update_user_password_by_id(
        self, db: Session, id: str, new_password: str
    ) -> bool:
        try:
            result = db.query(Auth).filter_by(id=id).update({"password": new_password})
            return True if result == 1 else False
        except:
            return False

    def update_email_by_id(self, db: Session, id: str, email: str) -> bool:
        try:
            result = db.query(Auth).filter_by(id=id).update({"email": email})
            return True if result == 1 else False
        except:
            return False

    def delete_auth_by_id(self, db: Session, id: str) -> bool:
        try:
            # Delete User
            result = Users.delete_user_by_id(db, id)

            if result:
                db.query(Auth).filter_by(id=id).delete()

                return True
            else:
                return False
        except:
            return False


Auths = AuthsTable()