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
        user = Auths.authenticate_user_by_trusted_header(trusted_email)
121
122

    if WEBUI_AUTH == False:
Timothy J. Baek's avatar
Timothy J. Baek committed
123
124
125
126

        if Users.get_num_users() != 0:
            raise HTTPException(400, detail=ERROR_MESSAGES.EXISTING_USERS)

127
128
129
        admin_email = "admin@localhost"
        admin_password = "admin"

Timothy J. Baek's avatar
Timothy J. Baek committed
130
        if not Users.get_user_by_email(admin_email.lower()):
131
132
133
134
135
            await signup(
                request,
                SignupForm(email=admin_email, password=admin_password, name="User"),
            )
        user = Auths.authenticate_user(admin_email.lower(), admin_password)
136
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
137
        user = Auths.authenticate_user(form_data.email.lower(), form_data.password)
138
139

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

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


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


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

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

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

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

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

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

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


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


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

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

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

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

270
271
272
273
274
275
276

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


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


@router.get("/signup/enabled/toggle", response_model=bool)
282
283
284
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
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307


############################
# 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
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
337


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


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