211 lines
		
	
	
		
			5.0 KiB
		
	
	
	
		
			Python
		
	
	
	
			
		
		
	
	
			211 lines
		
	
	
		
			5.0 KiB
		
	
	
	
		
			Python
		
	
	
	
| import logging
 | |
| import uuid
 | |
| from typing import Optional
 | |
| 
 | |
| from open_webui.internal.db import Base, get_db
 | |
| from open_webui.models.users import UserModel, Users
 | |
| from open_webui.env import SRC_LOG_LEVELS
 | |
| from pydantic import BaseModel
 | |
| from sqlalchemy import Boolean, Column, String, Text
 | |
| from open_webui.utils.auth import verify_password
 | |
| 
 | |
| 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 LdapForm(BaseModel):
 | |
|     user: 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}")
 | |
| 
 | |
|         user = Users.get_user_by_email(email)
 | |
|         if not user:
 | |
|             return None
 | |
| 
 | |
|         try:
 | |
|             with get_db() as db:
 | |
|                 auth = db.query(Auth).filter_by(id=user.id, active=True).first()
 | |
|                 if auth:
 | |
|                     if verify_password(password, auth.password):
 | |
|                         return user
 | |
|                     else:
 | |
|                         return None
 | |
|                 else:
 | |
|                     return None
 | |
|         except Exception:
 | |
|             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 Exception:
 | |
|             return False
 | |
| 
 | |
|     def authenticate_user_by_email(self, email: str) -> Optional[UserModel]:
 | |
|         log.info(f"authenticate_user_by_email: {email}")
 | |
|         try:
 | |
|             with get_db() as db:
 | |
|                 auth = db.query(Auth).filter_by(email=email, active=True).first()
 | |
|                 if auth:
 | |
|                     user = Users.get_user_by_id(auth.id)
 | |
|                     return user
 | |
|         except Exception:
 | |
|             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 Exception:
 | |
|             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 Exception:
 | |
|             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 Exception:
 | |
|             return False
 | |
| 
 | |
| 
 | |
| Auths = AuthsTable()
 |