auths.py 10.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
13
14
15

from apps.web.models.auths import (
    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
25
)
from apps.web.models.users import Users

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
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
113
        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
114
115
116
117
118
119
            await signup(
                request,
                SignupForm(
                    email=trusted_email, password=str(uuid.uuid4()), name=trusted_email
                ),
            )
120
121
        user = Auths.authenticate_user_by_trusted_header(trusted_email)
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
        if WEBUI_AUTH == False:
            if Users.get_num_users() != 0:
                raise HTTPException(400, detail=ERROR_MESSAGES.EXISTING_USERS)

            admin_email = "admin@localhost"
            admin_password = "admin"

            if not Users.get_user_by_email(admin_email.lower()):
                await signup(
                    request,
                    SignupForm(email=admin_email, password=admin_password, name="User"),
                )
            user = Auths.authenticate_user(admin_email.lower(), admin_password)
        else:
            user = Auths.authenticate_user(form_data.email.lower(), form_data.password)
137
138

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

        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
151
            "profile_image_url": user.profile_image_url,
152
153
        }
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
154
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
155
156
157
158
159
160
161
162


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


@router.post("/signup", response_model=SigninResponse)
163
async def signup(request: Request, form_data: SignupForm):
164
    if not request.app.state.ENABLE_SIGNUP:
Timothy J. Baek's avatar
Timothy J. Baek committed
165
166
167
        raise HTTPException(
            status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
        )
168

169
    if not validate_email_format(form_data.email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
170
171
172
        raise HTTPException(
            status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
        )
173

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

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

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

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

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


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


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

    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:
242
243

        print(form_data)
Timothy J. Baek's avatar
Timothy J. Baek committed
244
245
246
247
248
249
        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,
250
            form_data.role,
Timothy J. Baek's avatar
Timothy J. Baek committed
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
        )

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

269
270
271
272
273
274
275

############################
# ToggleSignUp
############################


@router.get("/signup/enabled", response_model=bool)
276
277
async def get_sign_up_status(request: Request, user=Depends(get_admin_user)):
    return request.app.state.ENABLE_SIGNUP
278
279
280


@router.get("/signup/enabled/toggle", response_model=bool)
281
282
283
async def toggle_sign_up(request: Request, user=Depends(get_admin_user)):
    request.app.state.ENABLE_SIGNUP = not request.app.state.ENABLE_SIGNUP
    return request.app.state.ENABLE_SIGNUP
Timothy J. Baek's avatar
Timothy J. Baek committed
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306


############################
# Default User Role
############################


@router.get("/signup/user/role")
async def get_default_user_role(request: Request, user=Depends(get_admin_user)):
    return request.app.state.DEFAULT_USER_ROLE


class UpdateRoleForm(BaseModel):
    role: str


@router.post("/signup/user/role")
async def update_default_user_role(
    request: Request, form_data: UpdateRoleForm, user=Depends(get_admin_user)
):
    if form_data.role in ["pending", "user", "admin"]:
        request.app.state.DEFAULT_USER_ROLE = form_data.role
    return request.app.state.DEFAULT_USER_ROLE
Timothy J. Baek's avatar
Timothy J. Baek committed
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336


############################
# JWT Expiration
############################


@router.get("/token/expires")
async def get_token_expires_duration(request: Request, user=Depends(get_admin_user)):
    return request.app.state.JWT_EXPIRES_IN


class UpdateJWTExpiresDurationForm(BaseModel):
    duration: str


@router.post("/token/expires/update")
async def update_token_expires_duration(
    request: Request,
    form_data: UpdateJWTExpiresDurationForm,
    user=Depends(get_admin_user),
):
    pattern = r"^(-1|0|(-?\d+(\.\d+)?)(ms|s|m|h|d|w))$"

    # Check if the input string matches the pattern
    if re.match(pattern, form_data.duration):
        request.app.state.JWT_EXPIRES_IN = form_data.duration
        return request.app.state.JWT_EXPIRES_IN
    else:
        return request.app.state.JWT_EXPIRES_IN
liu.vaayne's avatar
liu.vaayne committed
337
338
339
340
341
342
343
344
345
346
347


############################
# 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
348
    success = Users.update_user_api_key_by_id(user.id, api_key)
liu.vaayne's avatar
liu.vaayne committed
349
350
351
352
353
354
355
356
357
358
359
    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
360
    success = Users.update_user_api_key_by_id(user.id, None)
liu.vaayne's avatar
liu.vaayne committed
361
362
363
364
365
366
    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
367
    api_key = Users.get_user_api_key_by_id(user.id)
liu.vaayne's avatar
liu.vaayne committed
368
369
370
371
372
373
    if api_key:
        return {
            "api_key": api_key,
        }
    else:
        raise HTTPException(404, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)