auths.py 8.77 KB
Newer Older
liu.vaayne's avatar
liu.vaayne committed
1
2
from fastapi import Request
from fastapi import Depends, HTTPException, status
3

liu.vaayne's avatar
liu.vaayne committed
4
from fastapi import APIRouter
5
from pydantic import BaseModel
Timothy J. Baek's avatar
Timothy J. Baek committed
6
import re
7
8
9
10

from apps.web.models.auths import (
    SigninForm,
    SignupForm,
11
    UpdateProfileForm,
12
    UpdatePasswordForm,
13
14
15
    UserResponse,
    SigninResponse,
    Auths,
liu.vaayne's avatar
liu.vaayne committed
16
    ApiKey
17
18
19
)
from apps.web.models.users import Users

Timothy J. Baek's avatar
Timothy J. Baek committed
20
21
22
23
24
from utils.utils import (
    get_password_hash,
    get_current_user,
    get_admin_user,
    create_token,
liu.vaayne's avatar
liu.vaayne committed
25
    create_api_key
Timothy J. Baek's avatar
Timothy J. Baek committed
26
)
Timothy J. Baek's avatar
Timothy J. Baek committed
27
from utils.misc import parse_duration, validate_email_format
Timothy J. Baek's avatar
Timothy J. Baek committed
28
29
from utils.webhook import post_webhook
from constants import ERROR_MESSAGES, WEBHOOK_MESSAGES
30
from config import WEBUI_AUTH_TRUSTED_EMAIL_HEADER
31

Timothy J. Baek's avatar
Timothy J. Baek committed
32
33
router = APIRouter()

34
35
36
37
38
############################
# GetSessionUser
############################


39
@router.get("/", response_model=UserResponse)
40
41
42
43
44
45
46
47
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,
    }
48
49


50
############################
51
# Update Profile
52
53
54
55
############################


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


72
73
74
75
76
############################
# Update Password
############################


77
@router.post("/update/password", response_model=bool)
78
79
80
async def update_password(
    form_data: UpdatePasswordForm, session_user=Depends(get_current_user)
):
81
82
    if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
        raise HTTPException(400, detail=ERROR_MESSAGES.ACTION_PROHIBITED)
83
84
    if session_user:
        user = Auths.authenticate_user(session_user.email, form_data.password)
85

86
87
        if user:
            hashed = get_password_hash(form_data.new_password)
Timothy J. Baek's avatar
Timothy J. Baek committed
88
            return Auths.update_user_password_by_id(user.id, hashed)
89
90
        else:
            raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_PASSWORD)
91
92
93
94
    else:
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)


95
96
97
98
99
100
############################
# SignIn
############################


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

106
107
        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
108
109
110
111
112
113
            await signup(
                request,
                SignupForm(
                    email=trusted_email, password=str(uuid.uuid4()), name=trusted_email
                ),
            )
114
115
        user = Auths.authenticate_user_by_trusted_header(trusted_email)
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
116
        user = Auths.authenticate_user(form_data.email.lower(), form_data.password)
117
118

    if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
119
120
121
122
        token = create_token(
            data={"id": user.id},
            expires_delta=parse_duration(request.app.state.JWT_EXPIRES_IN),
        )
123
124
125
126
127
128
129
130

        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
131
            "profile_image_url": user.profile_image_url,
132
133
        }
    else:
Timothy J. Baek's avatar
Timothy J. Baek committed
134
        raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
135
136
137
138
139
140
141
142


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


@router.post("/signup", response_model=SigninResponse)
143
async def signup(request: Request, form_data: SignupForm):
144
    if not request.app.state.ENABLE_SIGNUP:
Timothy J. Baek's avatar
Timothy J. Baek committed
145
146
147
        raise HTTPException(
            status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
        )
148

149
    if not validate_email_format(form_data.email.lower()):
Timothy J. Baek's avatar
Timothy J. Baek committed
150
151
152
        raise HTTPException(
            status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.INVALID_EMAIL_FORMAT
        )
153

154
155
    if Users.get_user_by_email(form_data.email.lower()):
        raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)
156

157
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
158
159
160
161
162
        role = (
            "admin"
            if Users.get_num_users() == 0
            else request.app.state.DEFAULT_USER_ROLE
        )
163
        hashed = get_password_hash(form_data.password)
164
165
166
        user = Auths.insert_new_auth(
            form_data.email.lower(), hashed, form_data.name, role
        )
167

168
        if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
169
170
171
172
            token = create_token(
                data={"id": user.id},
                expires_delta=parse_duration(request.app.state.JWT_EXPIRES_IN),
            )
173
174
            # response.set_cookie(key='token', value=token, httponly=True)

Timothy J. Baek's avatar
Timothy J. Baek committed
175
176
177
            if request.app.state.WEBHOOK_URL:
                post_webhook(
                    request.app.state.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
178
                    WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
Timothy J. Baek's avatar
Timothy J. Baek committed
179
180
181
182
183
184
185
                    {
                        "action": "signup",
                        "message": WEBHOOK_MESSAGES.USER_SIGNUP(user.name),
                        "user": user.model_dump_json(exclude_none=True),
                    },
                )

186
187
188
189
190
191
192
193
194
195
            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:
196
            raise HTTPException(500, detail=ERROR_MESSAGES.CREATE_USER_ERROR)
197
    except Exception as err:
198
199
        raise HTTPException(500, detail=ERROR_MESSAGES.DEFAULT(err))

200
201
202
203
204
205
206

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


@router.get("/signup/enabled", response_model=bool)
207
208
async def get_sign_up_status(request: Request, user=Depends(get_admin_user)):
    return request.app.state.ENABLE_SIGNUP
209
210
211


@router.get("/signup/enabled/toggle", response_model=bool)
212
213
214
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
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237


############################
# 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
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267


############################
# 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
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304


############################
# 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()
    success = Auths.update_api_key_by_id(user.id, api_key)
    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)):
    success = Auths.update_api_key_by_id(user.id, None)
    return success


# get api key
@router.get("/api_key", response_model=ApiKey)
async def get_api_key(user=Depends(get_current_user)):
    api_key = Auths.get_api_key_by_id(user.id, None)
    if api_key:
        return {
            "api_key": api_key,
        }
    else:
        raise HTTPException(404, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)