auths.py 12.9 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.models.auths import (
14
15
    SigninForm,
    SignupForm,
Timothy J. Baek's avatar
Timothy J. Baek committed
16
    AddUserForm,
17
    UpdateProfileForm,
18
    UpdatePasswordForm,
19
20
21
    UserResponse,
    SigninResponse,
    Auths,
Timothy J. Baek's avatar
Timothy J. Baek committed
22
    ApiKey,
23
)
24
from apps.webui.models.users import Users
25

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

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

44
45
46
47
48
############################
# GetSessionUser
############################


49
@router.get("/", response_model=UserResponse)
Timothy J. Baek's avatar
Timothy J. Baek committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
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
    )

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


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


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


97
98
99
100
101
############################
# Update Password
############################


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

112
113
        if user:
            hashed = get_password_hash(form_data.new_password)
114
            return Auths.update_user_password_by_id(user.id, hashed)
115
116
        else:
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_PASSWORD)
117
118
119
120
    else:
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)


121
122
123
124
125
126
############################
# SignIn
############################


@router.post("/signin", response_model=SigninResponse)
127
async def signin(request: Request, response: Response, form_data: SigninForm):
128
129
    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
130
131
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_TRUSTED_HEADER)

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

150
151
        if Users.get_user_by_email(admin_email.lower()):
            user = Auths.authenticate_user(admin_email.lower(), admin_password)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
152
        else:
153
            if Users.get_num_users() != 0:
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
154
                raise HTTPException(400, detail=ERROR_MESSAGES.EXISTING_USERS)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
155

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
156
157
158
159
            await signup(
                request,
                SignupForm(email=admin_email, password=admin_password, name="User"),
            )
160

161
            user = Auths.authenticate_user(admin_email.lower(), admin_password)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
162
    else:
163
        user = Auths.authenticate_user(form_data.email.lower(), form_data.password)
164
165

    if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
166
167
        token = create_token(
            data={"id": user.id},
168
            expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
Timothy J. Baek's avatar
Timothy J. Baek committed
169
        )
170

Timothy J. Baek's avatar
Timothy J. Baek committed
171
172
173
174
175
176
177
        # Set the cookie token
        response.set_cookie(
            key="token",
            value=token,
            httponly=True,  # Ensures the cookie is not accessible via JavaScript
        )

178
179
180
181
182
183
184
        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
185
            "profile_image_url": user.profile_image_url,
186
187
        }
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
188
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
189
190
191
192
193
194
195
196


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


@router.post("/signup", response_model=SigninResponse)
197
async def signup(request: Request, response: Response, form_data: SignupForm):
198
    if not request.app.state.config.ENABLE_SIGNUP and WEBUI_AUTH:
Timothy J. Baek's avatar
Timothy J. Baek committed
199
200
201
        raise HTTPException(
            status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
        )
202

203
    if not validate_email_format(form_data.email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
204
205
206
        raise HTTPException(
            status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
        )
207

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

211
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
212
213
        role = (
            "admin"
214
            if Users.get_num_users() == 0
215
            else request.app.state.config.DEFAULT_USER_ROLE
Timothy J. Baek's avatar
Timothy J. Baek committed
216
        )
217
        hashed = get_password_hash(form_data.password)
218
        user = Auths.insert_new_auth(
Danny Liu's avatar
Danny Liu committed
219
220
221
222
223
            form_data.email.lower(),
            hashed,
            form_data.name,
            form_data.profile_image_url,
            role,
224
        )
225

226
        if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
227
228
            token = create_token(
                data={"id": user.id},
229
                expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
Timothy J. Baek's avatar
Timothy J. Baek committed
230
            )
231
232
            # response.set_cookie(key='token', value=token, httponly=True)

Timothy J. Baek's avatar
Timothy J. Baek committed
233
234
235
236
237
238
239
            # Set the cookie token
            response.set_cookie(
                key="token",
                value=token,
                httponly=True,  # Ensures the cookie is not accessible via JavaScript
            )

240
            if request.app.state.config.WEBHOOK_URL:
Timothy J. Baek's avatar
Timothy J. Baek committed
241
                post_webhook(
242
                    request.app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
243
                    WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
Timothy J. Baek's avatar
Timothy J. Baek committed
244
245
246
247
248
249
250
                    {
                        "action": "signup",
                        "message": WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
                        "user": user.model_dump_json(exclude_none=True),
                    },
                )

251
252
253
254
255
256
257
258
259
260
            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:
261
            raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
Timothy J. Baek's avatar
Timothy J. Baek committed
262
263
264
265
266
267
268
269
270
271
    except Exception as err:
        raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))


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


@router.post("/add", response_model=SigninResponse)
272
async def add_user(
273
    form_data: AddUserForm, user=Depends(get_admin_user)
274
):
Timothy J. Baek's avatar
Timothy J. Baek committed
275
276
277
278
279
280

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

281
    if Users.get_user_by_email(form_data.email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
282
283
284
        raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)

    try:
285
286

        print(form_data)
Timothy J. Baek's avatar
Timothy J. Baek committed
287
288
289
290
291
292
        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,
293
            form_data.role,
Timothy J. Baek's avatar
Timothy J. Baek committed
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
        )

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

312
313

############################
314
# GetAdminDetails
315
316
317
############################


318
@router.get("/admin/details")
319
async def get_admin_details(
320
    request: Request, user=Depends(get_current_user)
321
):
322
323
324
325
326
    if request.app.state.config.SHOW_ADMIN_DETAILS:
        admin_email = request.app.state.config.ADMIN_EMAIL
        admin_name = None

        print(admin_email, admin_name)
327

328
        if admin_email:
329
            admin = Users.get_user_by_email(admin_email)
330
331
332
            if admin:
                admin_name = admin.name
        else:
333
            admin = Users.get_first_user()
334
335
336
            if admin:
                admin_email = admin.email
                admin_name = admin.name
337

338
339
340
341
342
343
        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
344
345
346


############################
347
# ToggleSignUp
Timothy J. Baek's avatar
Timothy J. Baek committed
348
349
350
############################


351
352
353
354
355
356
357
358
359
@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
360
361


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


370
371
372
@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
373
):
374
375
    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
376

377
378
    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
379
380
381
382

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

    # Check if the input string matches the pattern
383
384
385
386
387
388
389
390
391
392
393
394
395
396
    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
397
398
399
400
401
402
403
404
405


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


# create api key
@router.post("/api_key", response_model=ApiKey)
406
async def create_api_key_(user=Depends(get_current_user)):
liu.vaayne's avatar
liu.vaayne committed
407
    api_key = create_api_key()
408
    success = Users.update_user_api_key_by_id(user.id, api_key)
liu.vaayne's avatar
liu.vaayne committed
409
410
411
412
413
414
415
416
417
418
    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)
419
420
async def delete_api_key(user=Depends(get_current_user)):
    success = Users.update_user_api_key_by_id(user.id, None)
liu.vaayne's avatar
liu.vaayne committed
421
422
423
424
425
    return success


# get api key
@router.get("/api_key", response_model=ApiKey)
426
427
async def get_api_key(user=Depends(get_current_user)):
    api_key = Users.get_user_api_key_by_id(user.id)
liu.vaayne's avatar
liu.vaayne committed
428
429
430
431
432
433
    if api_key:
        return {
            "api_key": api_key,
        }
    else:
        raise HTTPException(404, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)