open-webui/backend/apps/web/models/auths.py

196 lines
4.5 KiB
Python
Raw Normal View History

2023-11-19 01:47:12 +01:00
from pydantic import BaseModel
from typing import List, Union, Optional
import time
import uuid
import logging
2023-12-26 06:44:28 +01:00
from peewee import *
2023-11-19 01:47:12 +01:00
from apps.web.models.users import UserModel, Users
2024-02-01 21:53:13 +01:00
from utils.utils import verify_password
2023-11-19 01:47:12 +01:00
2023-12-26 06:44:28 +01:00
from apps.web.internal.db import DB
2023-11-19 01:47:12 +01:00
from config import SRC_LOG_LEVELS
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])
2023-11-19 01:47:12 +01:00
####################
# DB MODEL
####################
2023-12-26 06:44:28 +01:00
class Auth(Model):
id = CharField(unique=True)
email = CharField()
password = CharField()
active = BooleanField()
2024-03-26 11:22:17 +01:00
api_key = CharField(null=True, unique=True)
2023-12-26 06:44:28 +01:00
class Meta:
database = DB
2023-11-19 01:47:12 +01:00
class AuthModel(BaseModel):
id: str
email: str
password: str
active: bool = True
2024-03-26 11:22:17 +01:00
api_key: Optional[str] = None
2023-11-19 01:47:12 +01:00
####################
# Forms
####################
class Token(BaseModel):
token: str
token_type: str
2024-03-26 11:22:17 +01:00
class ApiKey(BaseModel):
api_key: Optional[str] = None
2023-11-19 01:47:12 +01:00
class UserResponse(BaseModel):
id: str
email: str
name: str
role: str
2023-11-19 06:41:43 +01:00
profile_image_url: str
2023-11-19 01:47:12 +01:00
class SigninResponse(Token, UserResponse):
pass
class SigninForm(BaseModel):
email: str
password: str
2024-01-27 05:27:45 +01:00
class ProfileImageUrlForm(BaseModel):
profile_image_url: str
class UpdateProfileForm(BaseModel):
profile_image_url: str
name: str
2023-12-29 09:12:30 +01:00
class UpdatePasswordForm(BaseModel):
password: str
new_password: str
2023-11-19 01:47:12 +01:00
class SignupForm(BaseModel):
name: str
email: str
password: str
class AuthsTable:
def __init__(self, db):
self.db = db
2023-12-26 06:44:28 +01:00
self.db.create_tables([Auth])
2023-11-19 01:47:12 +01:00
2024-01-06 05:59:56 +01:00
def insert_new_auth(
self, email: str, password: str, name: str, role: str = "pending"
) -> Optional[UserModel]:
log.info("insert_new_auth")
2023-11-19 01:47:12 +01:00
id = str(uuid.uuid4())
2024-01-06 05:59:56 +01:00
auth = AuthModel(
**{"id": id, "email": email, "password": password, "active": True}
)
2023-12-26 06:44:28 +01:00
result = Auth.create(**auth.model_dump())
2023-11-19 01:47:12 +01:00
user = Users.insert_new_user(id, name, email, role)
if result and user:
return user
else:
return None
2024-01-06 05:59:56 +01:00
def authenticate_user(self, email: str, password: str) -> Optional[UserModel]:
log.info(f"authenticate_user: {email}")
2023-12-26 08:43:21 +01:00
try:
auth = Auth.get(Auth.email == email, Auth.active == True)
if auth:
if verify_password(password, auth.password):
user = Users.get_user_by_id(auth.id)
return user
else:
return None
2023-11-19 01:47:12 +01:00
else:
return None
2023-12-26 08:43:21 +01:00
except:
2023-11-19 01:47:12 +01:00
return None
2024-03-26 11:22:17 +01:00
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:
auth = Auth.get(Auth.api_key == api_key, Auth.active == True)
if auth:
user = Users.get_user_by_id(auth.id)
return user
else:
return None
except:
return None
2023-12-29 09:29:18 +01:00
def update_user_password_by_id(self, id: str, new_password: str) -> bool:
2023-12-29 09:12:30 +01:00
try:
2023-12-29 09:29:18 +01:00
query = Auth.update(password=new_password).where(Auth.id == id)
result = query.execute()
2023-12-29 09:31:23 +01:00
return True if result == 1 else False
2023-12-29 09:12:30 +01:00
except:
return False
2024-01-06 05:59:56 +01:00
def update_email_by_id(self, id: str, email: str) -> bool:
try:
query = Auth.update(email=email).where(Auth.id == id)
result = query.execute()
return True if result == 1 else False
except:
return False
2024-03-26 11:22:17 +01:00
def update_api_key_by_id(self, id: str, api_key: str) -> str:
try:
query = Auth.update(api_key=api_key).where(Auth.id == id)
result = query.execute()
return True if result == 1 else False
except:
return False
def get_api_key_by_id(self, id: str) -> Optional[str]:
try:
auth = Auth.get(Auth.id == id)
return auth.api_key
except:
return None
2023-12-29 09:12:30 +01:00
def delete_auth_by_id(self, id: str) -> bool:
2023-12-29 08:24:51 +01:00
try:
# Delete User
result = Users.delete_user_by_id(id)
if result:
# Delete Auth
query = Auth.delete().where(Auth.id == id)
2024-01-06 05:59:56 +01:00
query.execute() # Remove the rows, return number of rows removed.
2023-12-29 08:24:51 +01:00
return True
else:
return False
except:
return False
2023-11-19 01:47:12 +01:00
Auths = AuthsTable(DB)