from pydantic import BaseModel from typing import Optional import uuid import logging from sqlalchemy import String, Column, Boolean, Text from apps.webui.models.users import UserModel, Users from utils.utils import verify_password from apps.webui.internal.db import Base, get_db 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(Text) 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, email: str, password: str, name: str, profile_image_url: str = "/user.png", role: str = "pending", oauth_sub: Optional[str] = None, ) -> Optional[UserModel]: with get_db() as db: 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( 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, email: str, password: str) -> Optional[UserModel]: log.info(f"authenticate_user: {email}") try: with get_db() as db: 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(auth.id) return user else: return None else: return None except: return None def authenticate_user_by_api_key(self, 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(api_key) return user if user else None except: return False def authenticate_user_by_trusted_header(self, email: str) -> Optional[UserModel]: log.info(f"authenticate_user_by_trusted_header: {email}") try: with get_db() as db: 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, id: str, new_password: str) -> bool: try: with get_db() as db: result = ( db.query(Auth).filter_by(id=id).update({"password": new_password}) ) db.commit() return True if result == 1 else False except: return False def update_email_by_id(self, id: str, email: str) -> bool: try: with get_db() as db: result = db.query(Auth).filter_by(id=id).update({"email": email}) db.commit() return True if result == 1 else False except: return False def delete_auth_by_id(self, id: str) -> bool: try: with get_db() as db: # Delete User result = Users.delete_user_by_id(id) if result: db.query(Auth).filter_by(id=id).delete() db.commit() return True else: return False except: return False Auths = AuthsTable()