auths.py 4.58 KB
Newer Older
1
from pydantic import BaseModel
2
from typing import Optional
3
import uuid
4
import logging
5
6
from sqlalchemy import String, Column, Boolean
from sqlalchemy.orm import Session
7

8
from apps.webui.models.users import UserModel, Users
Tim Farrell's avatar
Tim Farrell committed
9
from utils.utils import verify_password
10

11
from apps.webui.internal.db import Base
12

13
from config import SRC_LOG_LEVELS
Timothy J. Baek's avatar
Timothy J. Baek committed
14

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

18
19
20
21
22
####################
# DB MODEL
####################


23
24
class Auth(Base):
    __tablename__ = "auth"
Timothy J. Baek's avatar
Timothy J. Baek committed
25

26
27
28
29
    id = Column(String, primary_key=True)
    email = Column(String)
    password = Column(String)
    active = Column(Boolean)
Timothy J. Baek's avatar
Timothy J. Baek committed
30
31


32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class AuthModel(BaseModel):
    id: str
    email: str
    password: str
    active: bool = True


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


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

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
48

liu.vaayne's avatar
liu.vaayne committed
49
50
class ApiKey(BaseModel):
    api_key: Optional[str] = None
51

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
52

53
54
55
56
57
class UserResponse(BaseModel):
    id: str
    email: str
    name: str
    role: str
Timothy J. Baek's avatar
Timothy J. Baek committed
58
    profile_image_url: str
59
60
61
62
63
64
65
66
67
68
69


class SigninResponse(Token, UserResponse):
    pass


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


70
71
72
73
class ProfileImageUrlForm(BaseModel):
    profile_image_url: str


74
75
76
77
78
class UpdateProfileForm(BaseModel):
    profile_image_url: str
    name: str


79
80
81
82
83
class UpdatePasswordForm(BaseModel):
    password: str
    new_password: str


84
85
86
87
class SignupForm(BaseModel):
    name: str
    email: str
    password: str
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
88
    profile_image_url: Optional[str] = "/user.png"
89
90


Timothy J. Baek's avatar
Timothy J. Baek committed
91
class AddUserForm(SignupForm):
92
    role: Optional[str] = "pending"
Timothy J. Baek's avatar
Timothy J. Baek committed
93
94


95
96
class AuthsTable:

Timothy J. Baek's avatar
Timothy J. Baek committed
97
    def insert_new_auth(
Danny Liu's avatar
Danny Liu committed
98
        self,
99
        db: Session,
Danny Liu's avatar
Danny Liu committed
100
101
102
        email: str,
        password: str,
        name: str,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
103
        profile_image_url: str = "/user.png",
Danny Liu's avatar
Danny Liu committed
104
        role: str = "pending",
105
        oauth_sub: Optional[str] = None,
Timothy J. Baek's avatar
Timothy J. Baek committed
106
    ) -> Optional[UserModel]:
107
        log.info("insert_new_auth")
108
109
110

        id = str(uuid.uuid4())

Timothy J. Baek's avatar
Timothy J. Baek committed
111
112
113
        auth = AuthModel(
            **{"id": id, "email": email, "password": password, "active": True}
        )
114
115
        result = Auth(**auth.model_dump())
        db.add(result)
Timothy J. Baek's avatar
Timothy J. Baek committed
116

117
        user = Users.insert_new_user(
118
            db, id, name, email, profile_image_url, role, oauth_sub
119
        )
120

121
122
123
        db.commit()
        db.refresh(result)

124
125
126
127
128
        if result and user:
            return user
        else:
            return None

129
130
131
    def authenticate_user(
        self, db: Session, email: str, password: str
    ) -> Optional[UserModel]:
132
        log.info(f"authenticate_user: {email}")
Timothy J. Baek's avatar
Timothy J. Baek committed
133
        try:
134
            auth = db.query(Auth).filter_by(email=email, active=True).first()
Timothy J. Baek's avatar
Timothy J. Baek committed
135
136
            if auth:
                if verify_password(password, auth.password):
137
                    user = Users.get_user_by_id(db, auth.id)
Timothy J. Baek's avatar
Timothy J. Baek committed
138
139
140
                    return user
                else:
                    return None
141
142
            else:
                return None
Timothy J. Baek's avatar
Timothy J. Baek committed
143
        except:
144
145
            return None

146
147
148
    def authenticate_user_by_api_key(
        self, db: Session, api_key: str
    ) -> Optional[UserModel]:
liu.vaayne's avatar
liu.vaayne committed
149
150
151
152
        log.info(f"authenticate_user_by_api_key: {api_key}")
        # if no api_key, return None
        if not api_key:
            return None
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
153

liu.vaayne's avatar
liu.vaayne committed
154
        try:
155
            user = Users.get_user_by_api_key(db, api_key)
156
            return user if user else None
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
157
158
159
        except:
            return False

160
161
162
    def authenticate_user_by_trusted_header(
        self, db: Session, email: str
    ) -> Optional[UserModel]:
163
164
        log.info(f"authenticate_user_by_trusted_header: {email}")
        try:
165
            auth = db.query(Auth).filter(email=email, active=True).first()
166
167
168
            if auth:
                user = Users.get_user_by_id(auth.id)
                return user
liu.vaayne's avatar
liu.vaayne committed
169
170
171
        except:
            return None

172
173
174
    def update_user_password_by_id(
        self, db: Session, id: str, new_password: str
    ) -> bool:
175
        try:
176
            result = db.query(Auth).filter_by(id=id).update({"password": new_password})
Timothy J. Baek's avatar
Timothy J. Baek committed
177
            return True if result == 1 else False
178
179
180
        except:
            return False

181
    def update_email_by_id(self, db: Session, id: str, email: str) -> bool:
Timothy J. Baek's avatar
Timothy J. Baek committed
182
        try:
183
            result = db.query(Auth).filter_by(id=id).update({"email": email})
Timothy J. Baek's avatar
Timothy J. Baek committed
184
185
186
187
            return True if result == 1 else False
        except:
            return False

188
    def delete_auth_by_id(self, db: Session, id: str) -> bool:
Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
        try:
            # Delete User
191
            result = Users.delete_user_by_id(db, id)
Timothy J. Baek's avatar
Timothy J. Baek committed
192
193

            if result:
194
                db.query(Auth).filter_by(id=id).delete()
Timothy J. Baek's avatar
Timothy J. Baek committed
195
196
197
198
199
200
201

                return True
            else:
                return False
        except:
            return False

202

203
Auths = AuthsTable()