auths.py 12.5 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

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
Timothy J. Baek's avatar
Timothy J. Baek committed
37
38
39
40
41
from config import (
    WEBUI_AUTH,
    WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
    WEBUI_AUTH_TRUSTED_NAME_HEADER,
)
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
66
67
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
        secure=True,  # Ensures the cookie is sent over https
        samesite="lax",
    )

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


77
############################
78
# Update Profile
79
80
81
82
############################


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


99
100
101
102
103
############################
# Update Password
############################


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

113
114
        if user:
            hashed = get_password_hash(form_data.new_password)
Timothy J. Baek's avatar
Timothy J. Baek committed
115
            return Auths.update_user_password_by_id(user.id, hashed)
116
117
        else:
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_PASSWORD)
118
119
120
121
    else:
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)


122
123
124
125
126
127
############################
# SignIn
############################


@router.post("/signin", response_model=SigninResponse)
Timothy J. Baek's avatar
Timothy J. Baek committed
128
async def signin(request: Request, form_data: SigninForm):
129
130
    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
131
132
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_TRUSTED_HEADER)

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

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
151
152
153
154
155
        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
156

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

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

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

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


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


@router.post("/signup", response_model=SigninResponse)
191
async def signup(request: Request, form_data: SignupForm):
192
    if not request.app.state.config.ENABLE_SIGNUP and WEBUI_AUTH:
Timothy J. Baek's avatar
Timothy J. Baek committed
193
194
195
        raise HTTPException(
            status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
        )
196

197
    if not validate_email_format(form_data.email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
198
199
200
        raise HTTPException(
            status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
        )
201

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

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

220
        if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
221
222
            token = create_token(
                data={"id": user.id},
223
                expires_delta=parse_duration(request.app.state.config.JWT_EXPIRES_IN),
Timothy J. Baek's avatar
Timothy J. Baek committed
224
            )
225
226
            # response.set_cookie(key='token', value=token, httponly=True)

227
            if request.app.state.config.WEBHOOK_URL:
Timothy J. Baek's avatar
Timothy J. Baek committed
228
                post_webhook(
229
                    request.app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
230
                    WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
Timothy J. Baek's avatar
Timothy J. Baek committed
231
232
233
234
235
236
237
                    {
                        "action": "signup",
                        "message": WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
                        "user": user.model_dump_json(exclude_none=True),
                    },
                )

238
239
240
241
242
243
244
245
246
247
            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:
248
            raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
Timothy J. Baek's avatar
Timothy J. Baek committed
249
250
251
252
253
254
255
256
257
258
    except Exception as err:
        raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))


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


@router.post("/add", response_model=SigninResponse)
259
async def add_user(form_data: AddUserForm, user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
260
261
262
263
264
265
266
267
268
269

    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:
270
271

        print(form_data)
Timothy J. Baek's avatar
Timothy J. Baek committed
272
273
274
275
276
277
        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,
278
            form_data.role,
Timothy J. Baek's avatar
Timothy J. Baek committed
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
        )

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

297
298

############################
299
# GetAdminDetails
300
301
302
############################


303
304
305
306
307
308
309
@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)
310

311
312
313
314
315
316
317
318
319
        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
320

321
322
323
324
325
326
        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
327
328
329


############################
330
# ToggleSignUp
Timothy J. Baek's avatar
Timothy J. Baek committed
331
332
333
############################


334
335
336
337
338
339
340
341
342
@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
343
344


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


353
354
355
@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
356
):
357
358
    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
359

360
361
    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
362
363
364
365

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

    # Check if the input string matches the pattern
366
367
368
369
370
371
372
373
374
375
376
377
378
379
    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
380
381
382
383
384
385
386
387
388
389
390


############################
# 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
391
    success = Users.update_user_api_key_by_id(user.id, api_key)
liu.vaayne's avatar
liu.vaayne committed
392
393
394
395
396
397
398
399
400
401
402
    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
403
    success = Users.update_user_api_key_by_id(user.id, None)
liu.vaayne's avatar
liu.vaayne committed
404
405
406
407
408
409
    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
410
    api_key = Users.get_user_api_key_by_id(user.id)
liu.vaayne's avatar
liu.vaayne committed
411
412
413
414
415
416
    if api_key:
        return {
            "api_key": api_key,
        }
    else:
        raise HTTPException(404, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)