auths.py 11.8 KB
Newer Older
1
2
import logging

3
from fastapi import Request, UploadFile, File
liu.vaayne's avatar
liu.vaayne committed
4
from fastapi import Depends, HTTPException, status
5

liu.vaayne's avatar
liu.vaayne committed
6
from fastapi import APIRouter
7
from pydantic import BaseModel
Timothy J. Baek's avatar
Timothy J. Baek committed
8
import re
Timothy J. Baek's avatar
Timothy J. Baek committed
9
import uuid
10
11
import csv

12
from apps.webui.models.auths import (
13
14
    SigninForm,
    SignupForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
15
    AddUserForm,
16
    UpdateProfileForm,
17
    UpdatePasswordForm,
18
19
20
    UserResponse,
    SigninResponse,
    Auths,
Timothy J. Baek's avatar
Timothy J. Baek committed
21
    ApiKey,
22
)
23
from apps.webui.models.users import Users
24

Timothy J. Baek's avatar
Timothy J. Baek committed
25
26
27
28
29
from utils.utils import (
    get_password_hash,
    get_current_user,
    get_admin_user,
    create_token,
Timothy J. Baek's avatar
Timothy J. Baek committed
30
    create_api_key,
Timothy J. Baek's avatar
Timothy J. Baek committed
31
)
Timothy J. Baek's avatar
Timothy J. Baek committed
32
from utils.misc import parse_duration, validate_email_format
Timothy J. Baek's avatar
Timothy J. Baek committed
33
34
from utils.webhook import post_webhook
from constants import ERROR_MESSAGES, WEBHOOK_MESSAGES
35
36
37
38
from config import (
    WEBUI_AUTH,
    WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
)
39

Timothy J. Baek's avatar
Timothy J. Baek committed
40
41
router = APIRouter()

42
43
44
45
46
############################
# GetSessionUser
############################


47
@router.get("/", response_model=UserResponse)
48
49
50
51
52
53
54
55
async def get_session_user(user=Depends(get_current_user)):
    return {
        "id": user.id,
        "email": user.email,
        "name": user.name,
        "role": user.role,
        "profile_image_url": user.profile_image_url,
    }
56
57


58
############################
59
# Update Profile
60
61
62
63
############################


@router.post("/update/profile", response_model=UserResponse)
64
65
async def update_profile(
    form_data: UpdateProfileForm, session_user=Depends(get_current_user)
66
67
):
    if session_user:
68
69
70
        user = Users.update_user_by_id(
            session_user.id,
            {"profile_image_url": form_data.profile_image_url, "name": form_data.name},
71
72
73
74
75
76
77
78
79
        )
        if user:
            return user
        else:
            raise HTTPException(400, detail=ERROR_MESSAGES.DEFAULT())
    else:
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)


80
81
82
83
84
############################
# Update Password
############################


85
@router.post("/update/password", response_model=bool)
86
87
88
async def update_password(
    form_data: UpdatePasswordForm, session_user=Depends(get_current_user)
):
89
90
    if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
        raise HTTPException(400, detail=ERROR_MESSAGES.ACTION_PROHIBITED)
91
92
    if session_user:
        user = Auths.authenticate_user(session_user.email, form_data.password)
93

94
95
        if user:
            hashed = get_password_hash(form_data.new_password)
Timothy J. Baek's avatar
Timothy J. Baek committed
96
            return Auths.update_user_password_by_id(user.id, hashed)
97
98
        else:
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_PASSWORD)
99
100
101
102
    else:
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)


103
104
105
106
107
108
############################
# SignIn
############################


@router.post("/signin", response_model=SigninResponse)
Timothy J. Baek's avatar
Timothy J. Baek committed
109
async def signin(request: Request, form_data: SigninForm):
110
111
    if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
        if WEBUI_AUTH_TRUSTED_EMAIL_HEADER not in request.headers:
Timothy J. Baek's avatar
Timothy J. Baek committed
112
113
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_TRUSTED_HEADER)

114
115
        trusted_email = request.headers[WEBUI_AUTH_TRUSTED_EMAIL_HEADER].lower()
        if not Users.get_user_by_email(trusted_email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
116
117
118
119
120
121
            await signup(
                request,
                SignupForm(
                    email=trusted_email, password=str(uuid.uuid4()), name=trusted_email
                ),
            )
122
        user = Auths.authenticate_user_by_trusted_header(trusted_email)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
123
124
125
    elif WEBUI_AUTH == False:
        admin_email = "admin@localhost"
        admin_password = "admin"
Timothy J. Baek's avatar
Timothy J. Baek committed
126

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
127
128
129
130
131
        if Users.get_user_by_email(admin_email.lower()):
            user = Auths.authenticate_user(admin_email.lower(), admin_password)
        else:
            if Users.get_num_users() != 0:
                raise HTTPException(400, detail=ERROR_MESSAGES.EXISTING_USERS)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
132

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
133
134
135
136
            await signup(
                request,
                SignupForm(email=admin_email, password=admin_password, name="User"),
            )
137

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
138
139
140
            user = Auths.authenticate_user(admin_email.lower(), admin_password)
    else:
        user = Auths.authenticate_user(form_data.email.lower(), form_data.password)
141
142

    if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
143
144
        token = create_token(
            data={"id": user.id},
145
            expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
Timothy J. Baek's avatar
Timothy J. Baek committed
146
        )
147
148
149
150
151
152
153
154

        return {
            "token": token,
            "token_type": "Bearer",
            "id": user.id,
            "email": user.email,
            "name": user.name,
            "role": user.role,
Timothy J. Baek's avatar
Timothy J. Baek committed
155
            "profile_image_url": user.profile_image_url,
156
157
        }
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
158
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
159
160
161
162
163
164
165
166


############################
# SignUp
############################


@router.post("/signup", response_model=SigninResponse)
167
async def signup(request: Request, form_data: SignupForm):
168
    if not request.app.state.config.ENABLE_SIGNUP and WEBUI_AUTH:
Timothy J. Baek's avatar
Timothy J. Baek committed
169
170
171
        raise HTTPException(
            status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
        )
172

173
    if not validate_email_format(form_data.email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
174
175
176
        raise HTTPException(
            status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
        )
177

178
179
    if Users.get_user_by_email(form_data.email.lower()):
        raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)
180

181
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
182
183
184
        role = (
            "admin"
            if Users.get_num_users() == 0
185
            else request.app.state.config.DEFAULT_USER_ROLE
Timothy J. Baek's avatar
Timothy J. Baek committed
186
        )
187
        hashed = get_password_hash(form_data.password)
188
        user = Auths.insert_new_auth(
Danny Liu's avatar
Danny Liu committed
189
190
191
192
193
            form_data.email.lower(),
            hashed,
            form_data.name,
            form_data.profile_image_url,
            role,
194
        )
195

196
        if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
197
198
            token = create_token(
                data={"id": user.id},
199
                expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
Timothy J. Baek's avatar
Timothy J. Baek committed
200
            )
201
202
            # response.set_cookie(key='token', value=token, httponly=True)

203
            if request.app.state.config.WEBHOOK_URL:
Timothy J. Baek's avatar
Timothy J. Baek committed
204
                post_webhook(
205
                    request.app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
206
                    WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
Timothy J. Baek's avatar
Timothy J. Baek committed
207
208
209
210
211
212
213
                    {
                        "action": "signup",
                        "message": WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
                        "user": user.model_dump_json(exclude_none=True),
                    },
                )

214
215
216
217
218
219
220
221
222
223
            return {
                "token": token,
                "token_type": "Bearer",
                "id": user.id,
                "email": user.email,
                "name": user.name,
                "role": user.role,
                "profile_image_url": user.profile_image_url,
            }
        else:
224
            raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
Timothy J. Baek's avatar
Timothy J. Baek committed
225
226
227
228
229
230
231
232
233
234
    except Exception as err:
        raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))


############################
# AddUser
############################


@router.post("/add", response_model=SigninResponse)
235
async def add_user(form_data: AddUserForm, user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
236
237
238
239
240
241
242
243
244
245

    if not validate_email_format(form_data.email.lower()):
        raise HTTPException(
            status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
        )

    if Users.get_user_by_email(form_data.email.lower()):
        raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)

    try:
246
247

        print(form_data)
Timothy J. Baek's avatar
Timothy J. Baek committed
248
249
250
251
252
253
        hashed = get_password_hash(form_data.password)
        user = Auths.insert_new_auth(
            form_data.email.lower(),
            hashed,
            form_data.name,
            form_data.profile_image_url,
254
            form_data.role,
Timothy J. Baek's avatar
Timothy J. Baek committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
        )

        if user:
            token = create_token(data={"id": user.id})
            return {
                "token": token,
                "token_type": "Bearer",
                "id": user.id,
                "email": user.email,
                "name": user.name,
                "role": user.role,
                "profile_image_url": user.profile_image_url,
            }
        else:
            raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
270
    except Exception as err:
271
272
        raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))

273
274

############################
275
# GetAdminDetails
276
277
278
############################


279
280
281
282
283
284
285
@router.get("/admin/details")
async def get_admin_details(request: Request, user=Depends(get_current_user)):
    if request.app.state.config.SHOW_ADMIN_DETAILS:
        admin_email = request.app.state.config.ADMIN_EMAIL
        admin_name = None

        print(admin_email, admin_name)
286

287
288
289
290
291
292
293
294
295
        if admin_email:
            admin = Users.get_user_by_email(admin_email)
            if admin:
                admin_name = admin.name
        else:
            admin = Users.get_first_user()
            if admin:
                admin_email = admin.email
                admin_name = admin.name
296

297
298
299
300
301
302
        return {
            "name": admin_name,
            "email": admin_email,
        }
    else:
        raise HTTPException(400, detail=ERROR_MESSAGES.ACTION_PROHIBITED)
Timothy J. Baek's avatar
Timothy J. Baek committed
303
304
305


############################
306
# ToggleSignUp
Timothy J. Baek's avatar
Timothy J. Baek committed
307
308
309
############################


310
311
312
313
314
315
316
317
318
@router.get("/admin/config")
async def get_admin_config(request: Request, user=Depends(get_admin_user)):
    return {
        "SHOW_ADMIN_DETAILS": request.app.state.config.SHOW_ADMIN_DETAILS,
        "ENABLE_SIGNUP": request.app.state.config.ENABLE_SIGNUP,
        "DEFAULT_USER_ROLE": request.app.state.config.DEFAULT_USER_ROLE,
        "JWT_EXPIRES_IN": request.app.state.config.JWT_EXPIRES_IN,
        "ENABLE_COMMUNITY_SHARING": request.app.state.config.ENABLE_COMMUNITY_SHARING,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
319
320


321
322
323
324
325
326
class AdminConfig(BaseModel):
    SHOW_ADMIN_DETAILS: bool
    ENABLE_SIGNUP: bool
    DEFAULT_USER_ROLE: str
    JWT_EXPIRES_IN: str
    ENABLE_COMMUNITY_SHARING: bool
Timothy J. Baek's avatar
Timothy J. Baek committed
327
328


329
330
331
@router.post("/admin/config")
async def update_admin_config(
    request: Request, form_data: AdminConfig, user=Depends(get_admin_user)
Timothy J. Baek's avatar
Timothy J. Baek committed
332
):
333
334
    request.app.state.config.SHOW_ADMIN_DETAILS = form_data.SHOW_ADMIN_DETAILS
    request.app.state.config.ENABLE_SIGNUP = form_data.ENABLE_SIGNUP
Timothy J. Baek's avatar
Timothy J. Baek committed
335

336
337
    if form_data.DEFAULT_USER_ROLE in ["pending", "user", "admin"]:
        request.app.state.config.DEFAULT_USER_ROLE = form_data.DEFAULT_USER_ROLE
Timothy J. Baek's avatar
Timothy J. Baek committed
338
339
340
341

    pattern = r"^(-1|0|(-?\d+(\.\d+)?)(ms|s|m|h|d|w))$"

    # Check if the input string matches the pattern
342
343
344
345
346
347
348
349
350
351
352
353
354
355
    if re.match(pattern, form_data.JWT_EXPIRES_IN):
        request.app.state.config.JWT_EXPIRES_IN = form_data.JWT_EXPIRES_IN

    request.app.state.config.ENABLE_COMMUNITY_SHARING = (
        form_data.ENABLE_COMMUNITY_SHARING
    )

    return {
        "SHOW_ADMIN_DETAILS": request.app.state.config.SHOW_ADMIN_DETAILS,
        "ENABLE_SIGNUP": request.app.state.config.ENABLE_SIGNUP,
        "DEFAULT_USER_ROLE": request.app.state.config.DEFAULT_USER_ROLE,
        "JWT_EXPIRES_IN": request.app.state.config.JWT_EXPIRES_IN,
        "ENABLE_COMMUNITY_SHARING": request.app.state.config.ENABLE_COMMUNITY_SHARING,
    }
liu.vaayne's avatar
liu.vaayne committed
356
357
358
359
360
361
362
363
364
365
366


############################
# API Key
############################


# create api key
@router.post("/api_key", response_model=ApiKey)
async def create_api_key_(user=Depends(get_current_user)):
    api_key = create_api_key()
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
367
    success = Users.update_user_api_key_by_id(user.id, api_key)
liu.vaayne's avatar
liu.vaayne committed
368
369
370
371
372
373
374
375
376
377
378
    if success:
        return {
            "api_key": api_key,
        }
    else:
        raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_API_KEY_ERROR)


# delete api key
@router.delete("/api_key", response_model=bool)
async def delete_api_key(user=Depends(get_current_user)):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
379
    success = Users.update_user_api_key_by_id(user.id, None)
liu.vaayne's avatar
liu.vaayne committed
380
381
382
383
384
385
    return success


# get api key
@router.get("/api_key", response_model=ApiKey)
async def get_api_key(user=Depends(get_current_user)):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
386
    api_key = Users.get_user_api_key_by_id(user.id)
liu.vaayne's avatar
liu.vaayne committed
387
388
389
390
391
392
    if api_key:
        return {
            "api_key": api_key,
        }
    else:
        raise HTTPException(404, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)