auths.py 13.2 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
Timothy J. Baek's avatar
Timothy J. Baek committed
5
from fastapi.responses import Response
6

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

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
43
44
router = APIRouter()

45
46
47
48
49
############################
# GetSessionUser
############################


50
@router.get("/", response_model=UserResponse)
Timothy J. Baek's avatar
Timothy J. Baek committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
async def get_session_user(
    request: Request, response: Response, user=Depends(get_current_user)
):
    token = create_token(
        data={"id": user.id},
        expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
    )

    # Set the cookie token
    response.set_cookie(
        key="token",
        value=token,
        httponly=True,  # Ensures the cookie is not accessible via JavaScript
    )

66
67
68
69
70
71
72
    return {
        "id": user.id,
        "email": user.email,
        "name": user.name,
        "role": user.role,
        "profile_image_url": user.profile_image_url,
    }
73
74


75
############################
76
# Update Profile
77
78
79
80
############################


@router.post("/update/profile", response_model=UserResponse)
81
async def update_profile(
82
83
84
    form_data: UpdateProfileForm,
    session_user=Depends(get_current_user),
    db=Depends(get_db),
85
86
):
    if session_user:
87
        user = Users.update_user_by_id(
88
            db,
89
90
            session_user.id,
            {"profile_image_url": form_data.profile_image_url, "name": form_data.name},
91
92
93
94
95
96
97
98
99
        )
        if user:
            return user
        else:
            raise HTTPException(400, detail=ERROR_MESSAGES.DEFAULT())
    else:
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)


100
101
102
103
104
############################
# Update Password
############################


105
@router.post("/update/password", response_model=bool)
106
async def update_password(
107
108
109
    form_data: UpdatePasswordForm,
    session_user=Depends(get_current_user),
    db=Depends(get_db),
110
):
111
112
    if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
        raise HTTPException(400, detail=ERROR_MESSAGES.ACTION_PROHIBITED)
113
    if session_user:
114
        user = Auths.authenticate_user(db, session_user.email, form_data.password)
115

116
117
        if user:
            hashed = get_password_hash(form_data.new_password)
118
            return Auths.update_user_password_by_id(db, user.id, hashed)
119
120
        else:
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_PASSWORD)
121
122
123
124
    else:
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)


125
126
127
128
129
130
############################
# SignIn
############################


@router.post("/signin", response_model=SigninResponse)
131
async def signin(request: Request, response: Response, form_data: SigninForm, db=Depends(get_db)):
132
133
    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
134
135
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_TRUSTED_HEADER)

136
        trusted_email = request.headers[WEBUI_AUTH_TRUSTED_EMAIL_HEADER].lower()
137
138
        trusted_name = trusted_email
        if WEBUI_AUTH_TRUSTED_NAME_HEADER:
Timothy J. Baek's avatar
Timothy J. Baek committed
139
140
141
            trusted_name = request.headers.get(
                WEBUI_AUTH_TRUSTED_NAME_HEADER, trusted_email
            )
142
        if not Users.get_user_by_email(db, trusted_email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
143
144
145
            await signup(
                request,
                SignupForm(
146
                    email=trusted_email, password=str(uuid.uuid4()), name=trusted_name
Timothy J. Baek's avatar
Timothy J. Baek committed
147
                ),
148
                db,
Timothy J. Baek's avatar
Timothy J. Baek committed
149
            )
150
        user = Auths.authenticate_user_by_trusted_header(db, trusted_email)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
151
152
153
    elif WEBUI_AUTH == False:
        admin_email = "admin@localhost"
        admin_password = "admin"
Timothy J. Baek's avatar
Timothy J. Baek committed
154

155
156
        if Users.get_user_by_email(db, admin_email.lower()):
            user = Auths.authenticate_user(db, admin_email.lower(), admin_password)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
157
        else:
158
            if Users.get_num_users(db) != 0:
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
159
                raise HTTPException(400, detail=ERROR_MESSAGES.EXISTING_USERS)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
160

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
161
162
163
            await signup(
                request,
                SignupForm(email=admin_email, password=admin_password, name="User"),
164
                db,
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
165
            )
166

167
            user = Auths.authenticate_user(db, admin_email.lower(), admin_password)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
168
    else:
169
        user = Auths.authenticate_user(db, form_data.email.lower(), form_data.password)
170
171

    if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
172
173
        token = create_token(
            data={"id": user.id},
174
            expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
Timothy J. Baek's avatar
Timothy J. Baek committed
175
        )
176

Timothy J. Baek's avatar
Timothy J. Baek committed
177
178
179
180
181
182
183
        # Set the cookie token
        response.set_cookie(
            key="token",
            value=token,
            httponly=True,  # Ensures the cookie is not accessible via JavaScript
        )

184
185
186
187
188
189
190
        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
191
            "profile_image_url": user.profile_image_url,
192
193
        }
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
194
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
195
196
197
198
199
200
201
202


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


@router.post("/signup", response_model=SigninResponse)
203
async def signup(request: Request, response: Response, form_data: SignupForm, db=Depends(get_db)):
204
    if not request.app.state.config.ENABLE_SIGNUP and WEBUI_AUTH:
Timothy J. Baek's avatar
Timothy J. Baek committed
205
206
207
        raise HTTPException(
            status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
        )
208

209
    if not validate_email_format(form_data.email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
210
211
212
        raise HTTPException(
            status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
        )
213

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

217
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
218
219
        role = (
            "admin"
220
            if Users.get_num_users(db) == 0
221
            else request.app.state.config.DEFAULT_USER_ROLE
Timothy J. Baek's avatar
Timothy J. Baek committed
222
        )
223
        hashed = get_password_hash(form_data.password)
224
        user = Auths.insert_new_auth(
225
            db,
Danny Liu's avatar
Danny Liu committed
226
227
228
229
230
            form_data.email.lower(),
            hashed,
            form_data.name,
            form_data.profile_image_url,
            role,
231
        )
232

233
        if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
234
235
            token = create_token(
                data={"id": user.id},
236
                expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
Timothy J. Baek's avatar
Timothy J. Baek committed
237
            )
238
239
            # response.set_cookie(key='token', value=token, httponly=True)

Timothy J. Baek's avatar
Timothy J. Baek committed
240
241
242
243
244
245
246
            # Set the cookie token
            response.set_cookie(
                key="token",
                value=token,
                httponly=True,  # Ensures the cookie is not accessible via JavaScript
            )

247
            if request.app.state.config.WEBHOOK_URL:
Timothy J. Baek's avatar
Timothy J. Baek committed
248
                post_webhook(
249
                    request.app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
250
                    WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
Timothy J. Baek's avatar
Timothy J. Baek committed
251
252
253
254
255
256
257
                    {
                        "action": "signup",
                        "message": WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
                        "user": user.model_dump_json(exclude_none=True),
                    },
                )

258
259
260
261
262
263
264
265
266
267
            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:
268
            raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
Timothy J. Baek's avatar
Timothy J. Baek committed
269
270
271
272
273
274
275
276
277
278
    except Exception as err:
        raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))


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


@router.post("/add", response_model=SigninResponse)
279
280
281
async def add_user(
    form_data: AddUserForm, user=Depends(get_admin_user), db=Depends(get_db)
):
Timothy J. Baek's avatar
Timothy J. Baek committed
282
283
284
285
286
287

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

288
    if Users.get_user_by_email(db, form_data.email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
289
290
291
        raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)

    try:
292
293

        print(form_data)
Timothy J. Baek's avatar
Timothy J. Baek committed
294
295
        hashed = get_password_hash(form_data.password)
        user = Auths.insert_new_auth(
296
            db,
Timothy J. Baek's avatar
Timothy J. Baek committed
297
298
299
300
            form_data.email.lower(),
            hashed,
            form_data.name,
            form_data.profile_image_url,
301
            form_data.role,
Timothy J. Baek's avatar
Timothy J. Baek committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
        )

        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)
317
    except Exception as err:
318
319
        raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))

320
321

############################
322
# GetAdminDetails
323
324
325
############################


326
@router.get("/admin/details")
327
328
329
async def get_admin_details(
    request: Request, user=Depends(get_current_user), db=Depends(get_db)
):
330
331
332
333
334
    if request.app.state.config.SHOW_ADMIN_DETAILS:
        admin_email = request.app.state.config.ADMIN_EMAIL
        admin_name = None

        print(admin_email, admin_name)
335

336
        if admin_email:
337
            admin = Users.get_user_by_email(db, admin_email)
338
339
340
            if admin:
                admin_name = admin.name
        else:
341
            admin = Users.get_first_user(db)
342
343
344
            if admin:
                admin_email = admin.email
                admin_name = admin.name
345

346
347
348
349
350
351
        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
352
353
354


############################
355
# ToggleSignUp
Timothy J. Baek's avatar
Timothy J. Baek committed
356
357
358
############################


359
360
361
362
363
364
365
366
367
@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
368
369


370
371
372
373
374
375
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
376
377


378
379
380
@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
381
):
382
383
    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
384

385
386
    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
387
388
389
390

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

    # Check if the input string matches the pattern
391
392
393
394
395
396
397
398
399
400
401
402
403
404
    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
405
406
407
408
409
410
411
412
413


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


# create api key
@router.post("/api_key", response_model=ApiKey)
414
async def create_api_key_(user=Depends(get_current_user), db=Depends(get_db)):
liu.vaayne's avatar
liu.vaayne committed
415
    api_key = create_api_key()
416
    success = Users.update_user_api_key_by_id(db, user.id, api_key)
liu.vaayne's avatar
liu.vaayne committed
417
418
419
420
421
422
423
424
425
426
    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)
427
428
async def delete_api_key(user=Depends(get_current_user), db=Depends(get_db)):
    success = Users.update_user_api_key_by_id(db, user.id, None)
liu.vaayne's avatar
liu.vaayne committed
429
430
431
432
433
    return success


# get api key
@router.get("/api_key", response_model=ApiKey)
434
435
async def get_api_key(user=Depends(get_current_user), db=Depends(get_db)):
    api_key = Users.get_user_api_key_by_id(db, user.id)
liu.vaayne's avatar
liu.vaayne committed
436
437
438
439
440
441
    if api_key:
        return {
            "api_key": api_key,
        }
    else:
        raise HTTPException(404, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)