auths.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206
  1. import logging
  2. import uuid
  3. from typing import Optional
  4. from open_webui.internal.db import Base, get_db
  5. from open_webui.models.users import UserModel, Users
  6. from open_webui.env import SRC_LOG_LEVELS
  7. from pydantic import BaseModel
  8. from sqlalchemy import Boolean, Column, String, Text
  9. from open_webui.utils.auth import verify_password
  10. log = logging.getLogger(__name__)
  11. log.setLevel(SRC_LOG_LEVELS["MODELS"])
  12. ####################
  13. # DB MODEL
  14. ####################
  15. class Auth(Base):
  16. __tablename__ = "auth"
  17. id = Column(String, primary_key=True)
  18. email = Column(String)
  19. password = Column(Text)
  20. active = Column(Boolean)
  21. class AuthModel(BaseModel):
  22. id: str
  23. email: str
  24. password: str
  25. active: bool = True
  26. ####################
  27. # Forms
  28. ####################
  29. class Token(BaseModel):
  30. token: str
  31. token_type: str
  32. class ApiKey(BaseModel):
  33. api_key: Optional[str] = None
  34. class UserResponse(BaseModel):
  35. id: str
  36. email: str
  37. name: str
  38. role: str
  39. profile_image_url: str
  40. class SigninResponse(Token, UserResponse):
  41. pass
  42. class SigninForm(BaseModel):
  43. email: str
  44. password: str
  45. class LdapForm(BaseModel):
  46. user: str
  47. password: str
  48. class ProfileImageUrlForm(BaseModel):
  49. profile_image_url: str
  50. class UpdateProfileForm(BaseModel):
  51. profile_image_url: str
  52. name: str
  53. class UpdatePasswordForm(BaseModel):
  54. password: str
  55. new_password: str
  56. class SignupForm(BaseModel):
  57. name: str
  58. email: str
  59. password: str
  60. profile_image_url: Optional[str] = "/user.png"
  61. class AddUserForm(SignupForm):
  62. role: Optional[str] = "pending"
  63. class AuthsTable:
  64. def insert_new_auth(
  65. self,
  66. email: str,
  67. password: str,
  68. name: str,
  69. profile_image_url: str = "/user.png",
  70. role: str = "pending",
  71. oauth_sub: Optional[str] = None,
  72. ) -> Optional[UserModel]:
  73. with get_db() as db:
  74. log.info("insert_new_auth")
  75. id = str(uuid.uuid4())
  76. auth = AuthModel(
  77. **{"id": id, "email": email, "password": password, "active": True}
  78. )
  79. result = Auth(**auth.model_dump())
  80. db.add(result)
  81. user = Users.insert_new_user(
  82. id, name, email, profile_image_url, role, oauth_sub
  83. )
  84. db.commit()
  85. db.refresh(result)
  86. if result and user:
  87. return user
  88. else:
  89. return None
  90. def authenticate_user(self, email: str, password: str) -> Optional[UserModel]:
  91. log.info(f"authenticate_user: {email}")
  92. try:
  93. with get_db() as db:
  94. auth = db.query(Auth).filter_by(email=email, active=True).first()
  95. if auth:
  96. if verify_password(password, auth.password):
  97. user = Users.get_user_by_id(auth.id)
  98. return user
  99. else:
  100. return None
  101. else:
  102. return None
  103. except Exception:
  104. return None
  105. def authenticate_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
  106. log.info(f"authenticate_user_by_api_key: {api_key}")
  107. # if no api_key, return None
  108. if not api_key:
  109. return None
  110. try:
  111. user = Users.get_user_by_api_key(api_key)
  112. return user if user else None
  113. except Exception:
  114. return False
  115. def authenticate_user_by_trusted_header(self, email: str) -> Optional[UserModel]:
  116. log.info(f"authenticate_user_by_trusted_header: {email}")
  117. try:
  118. with get_db() as db:
  119. auth = db.query(Auth).filter_by(email=email, active=True).first()
  120. if auth:
  121. user = Users.get_user_by_id(auth.id)
  122. return user
  123. except Exception:
  124. return None
  125. def update_user_password_by_id(self, id: str, new_password: str) -> bool:
  126. try:
  127. with get_db() as db:
  128. result = (
  129. db.query(Auth).filter_by(id=id).update({"password": new_password})
  130. )
  131. db.commit()
  132. return True if result == 1 else False
  133. except Exception:
  134. return False
  135. def update_email_by_id(self, id: str, email: str) -> bool:
  136. try:
  137. with get_db() as db:
  138. result = db.query(Auth).filter_by(id=id).update({"email": email})
  139. db.commit()
  140. return True if result == 1 else False
  141. except Exception:
  142. return False
  143. def delete_auth_by_id(self, id: str) -> bool:
  144. try:
  145. with get_db() as db:
  146. # Delete User
  147. result = Users.delete_user_by_id(id)
  148. if result:
  149. db.query(Auth).filter_by(id=id).delete()
  150. db.commit()
  151. return True
  152. else:
  153. return False
  154. except Exception:
  155. return False
  156. Auths = AuthsTable()