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

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
from config import WEBUI_AUTH, WEBUI_AUTH_TRUSTED_EMAIL_HEADER, WEBUI_AUTH_TRUSTED_NAME_HEADER
37

Timothy J. Baek's avatar
Timothy J. Baek committed
38
39
router = APIRouter()

40
41
42
43
44
############################
# GetSessionUser
############################


45
@router.get("/", response_model=UserResponse)
46
47
48
49
50
51
52
53
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,
    }
54
55


56
############################
57
# Update Profile
58
59
60
61
############################


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


78
79
80
81
82
############################
# Update Password
############################


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

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


101
102
103
104
105
106
############################
# SignIn
############################


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

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

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
128
129
130
131
132
        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
133

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

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

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

        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
156
            "profile_image_url": user.profile_image_url,
157
158
        }
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
159
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
160
161
162
163
164
165
166
167


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


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

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

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

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

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

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

215
216
217
218
219
220
221
222
223
224
            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:
225
            raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
Timothy J. Baek's avatar
Timothy J. Baek committed
226
227
228
229
230
231
232
233
234
235
    except Exception as err:
        raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))


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


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

    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:
247
248

        print(form_data)
Timothy J. Baek's avatar
Timothy J. Baek committed
249
250
251
252
253
254
        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,
255
            form_data.role,
Timothy J. Baek's avatar
Timothy J. Baek committed
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
        )

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

274
275

############################
276
# GetAdminDetails
277
278
279
############################


280
281
282
283
284
285
286
@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)
287

288
289
290
291
292
293
294
295
296
        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
297

298
299
300
301
302
303
        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
304
305
306


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


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


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


330
331
332
@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
333
):
334
335
    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
336

337
338
    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
339
340
341
342

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

    # Check if the input string matches the pattern
343
344
345
346
347
348
349
350
351
352
353
354
355
356
    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
357
358
359
360
361
362
363
364
365
366
367


############################
# 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
368
    success = Users.update_user_api_key_by_id(user.id, api_key)
liu.vaayne's avatar
liu.vaayne committed
369
370
371
372
373
374
375
376
377
378
379
    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
380
    success = Users.update_user_api_key_by_id(user.id, None)
liu.vaayne's avatar
liu.vaayne committed
381
382
383
384
385
386
    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
387
    api_key = Users.get_user_api_key_by_id(user.id)
liu.vaayne's avatar
liu.vaayne committed
388
389
390
391
392
393
    if api_key:
        return {
            "api_key": api_key,
        }
    else:
        raise HTTPException(404, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)