1
- from typing import Any , Union , Dict
1
+ from typing import Any , Union
2
2
3
3
from fastapi import APIRouter , Body , Depends , HTTPException
4
4
from fastapi .security import OAuth2PasswordRequestForm
5
- from sqlalchemy . orm import Session
5
+ from motor . core import AgnosticDatabase
6
6
7
7
from app import crud , models , schemas
8
8
from app .api import deps
24
24
- The user ID or password was incorrect.
25
25
- The account does not exist.
26
26
- The account is locked or disabled.
27
- - Code should go through the same process, no matter what, allowing the application to return in approximately
27
+ - Code should go through the same process, no matter what, allowing the application to return in approximately
28
28
the same response time.
29
29
- In the words of George Orwell, break these rules sooner than do something truly barbaric.
30
30
33
33
34
34
35
35
@router .post ("/magic/{email}" , response_model = schemas .WebToken )
36
- def login_with_magic_link (* , db : Session = Depends (deps .get_db ), email : str ) -> Any :
36
+ async def login_with_magic_link (* , db : AgnosticDatabase = Depends (deps .get_db ), email : str ) -> Any :
37
37
"""
38
38
First step of a 'magic link' login. Check if the user exists and generate a magic link. Generates two short-duration
39
39
jwt tokens, one for validation, one for email. Creates user if not exist.
40
40
"""
41
- user = crud .user .get_by_email (db , email = email )
41
+ user = await crud .user .get_by_email (db , email = email )
42
42
if not user :
43
43
user_in = schemas .UserCreate (** {"email" : email })
44
- user = crud .user .create (db , obj_in = user_in )
44
+ user = await crud .user .create (db , obj_in = user_in )
45
45
if not crud .user .is_active (user ):
46
46
# Still permits a timed-attack, but does create ambiguity.
47
47
raise HTTPException (status_code = 400 , detail = "A link to activate your account has been emailed." )
@@ -53,9 +53,9 @@ def login_with_magic_link(*, db: Session = Depends(deps.get_db), email: str) ->
53
53
54
54
55
55
@router .post ("/claim" , response_model = schemas .Token )
56
- def validate_magic_link (
56
+ async def validate_magic_link (
57
57
* ,
58
- db : Session = Depends (deps .get_db ),
58
+ db : AgnosticDatabase = Depends (deps .get_db ),
59
59
obj_in : schemas .WebToken ,
60
60
magic_in : bool = Depends (deps .get_magic_token ),
61
61
) -> Any :
@@ -64,7 +64,7 @@ def validate_magic_link(
64
64
"""
65
65
claim_in = deps .get_magic_token (token = obj_in .claim )
66
66
# Get the user
67
- user = crud .user .get (db , id = magic_in .sub )
67
+ user = await crud .user .get (db , id = magic_in .sub )
68
68
# Test the claims
69
69
if (
70
70
(claim_in .sub == magic_in .sub )
@@ -75,15 +75,15 @@ def validate_magic_link(
75
75
raise HTTPException (status_code = 400 , detail = "Login failed; invalid claim." )
76
76
# Validate that the email is the user's
77
77
if not user .email_validated :
78
- crud .user .validate_email (db = db , db_obj = user )
78
+ await crud .user .validate_email (db = db , db_obj = user )
79
79
# Check if totp active
80
80
refresh_token = None
81
81
force_totp = True
82
82
if not user .totp_secret :
83
83
# No TOTP, so this concludes the login validation
84
84
force_totp = False
85
85
refresh_token = security .create_refresh_token (subject = user .id )
86
- crud .token .create (db = db , obj_in = refresh_token , user_obj = user )
86
+ await crud .token .create (db = db , obj_in = refresh_token , user_obj = user )
87
87
return {
88
88
"access_token" : security .create_access_token (subject = user .id , force_totp = force_totp ),
89
89
"refresh_token" : refresh_token ,
@@ -92,11 +92,13 @@ def validate_magic_link(
92
92
93
93
94
94
@router .post ("/oauth" , response_model = schemas .Token )
95
- def login_with_oauth2 (db : Session = Depends (deps .get_db ), form_data : OAuth2PasswordRequestForm = Depends ()) -> Any :
95
+ async def login_with_oauth2 (
96
+ db : AgnosticDatabase = Depends (deps .get_db ), form_data : OAuth2PasswordRequestForm = Depends ()
97
+ ) -> Any :
96
98
"""
97
99
First step with OAuth2 compatible token login, get an access token for future requests.
98
100
"""
99
- user = crud .user .authenticate (db , email = form_data .username , password = form_data .password )
101
+ user = await crud .user .authenticate (db , email = form_data .username , password = form_data .password )
100
102
if not form_data .password or not user or not crud .user .is_active (user ):
101
103
raise HTTPException (status_code = 400 , detail = "Login failed; incorrect email or password" )
102
104
# Check if totp active
@@ -106,7 +108,7 @@ def login_with_oauth2(db: Session = Depends(deps.get_db), form_data: OAuth2Passw
106
108
# No TOTP, so this concludes the login validation
107
109
force_totp = False
108
110
refresh_token = security .create_refresh_token (subject = user .id )
109
- crud .token .create (db = db , obj_in = refresh_token , user_obj = user )
111
+ await crud .token .create (db = db , obj_in = refresh_token , user_obj = user )
110
112
return {
111
113
"access_token" : security .create_access_token (subject = user .id , force_totp = force_totp ),
112
114
"refresh_token" : refresh_token ,
@@ -115,9 +117,9 @@ def login_with_oauth2(db: Session = Depends(deps.get_db), form_data: OAuth2Passw
115
117
116
118
117
119
@router .post ("/totp" , response_model = schemas .Token )
118
- def login_with_totp (
120
+ async def login_with_totp (
119
121
* ,
120
- db : Session = Depends (deps .get_db ),
122
+ db : AgnosticDatabase = Depends (deps .get_db ),
121
123
totp_data : schemas .WebToken ,
122
124
current_user : models .User = Depends (deps .get_totp_user ),
123
125
) -> Any :
@@ -130,9 +132,9 @@ def login_with_totp(
130
132
if not new_counter :
131
133
raise HTTPException (status_code = 400 , detail = "Login failed; unable to verify TOTP." )
132
134
# Save the new counter to prevent reuse
133
- current_user = crud .user .update_totp_counter (db = db , db_obj = current_user , new_counter = new_counter )
135
+ current_user = await crud .user .update_totp_counter (db = db , db_obj = current_user , new_counter = new_counter )
134
136
refresh_token = security .create_refresh_token (subject = current_user .id )
135
- crud .token .create (db = db , obj_in = refresh_token , user_obj = current_user )
137
+ await crud .token .create (db = db , obj_in = refresh_token , user_obj = current_user )
136
138
return {
137
139
"access_token" : security .create_access_token (subject = current_user .id ),
138
140
"refresh_token" : refresh_token ,
@@ -141,17 +143,17 @@ def login_with_totp(
141
143
142
144
143
145
@router .put ("/totp" , response_model = schemas .Msg )
144
- def enable_totp_authentication (
146
+ async def enable_totp_authentication (
145
147
* ,
146
- db : Session = Depends (deps .get_db ),
148
+ db : AgnosticDatabase = Depends (deps .get_db ),
147
149
data_in : schemas .EnableTOTP ,
148
150
current_user : models .User = Depends (deps .get_current_active_user ),
149
151
) -> Any :
150
152
"""
151
153
For validation of token before enabling TOTP.
152
154
"""
153
155
if current_user .hashed_password :
154
- user = crud .user .authenticate (db , email = current_user .email , password = data_in .password )
156
+ user = await crud .user .authenticate (db , email = current_user .email , password = data_in .password )
155
157
if not data_in .password or not user :
156
158
raise HTTPException (status_code = 400 , detail = "Unable to authenticate or activate TOTP." )
157
159
totp_in = security .create_new_totp (label = current_user .email , uri = data_in .uri )
@@ -161,39 +163,39 @@ def enable_totp_authentication(
161
163
if not new_counter :
162
164
raise HTTPException (status_code = 400 , detail = "Unable to authenticate or activate TOTP." )
163
165
# Enable TOTP and save the new counter to prevent reuse
164
- current_user = crud .user .activate_totp (db = db , db_obj = current_user , totp_in = totp_in )
165
- current_user = crud .user .update_totp_counter (db = db , db_obj = current_user , new_counter = new_counter )
166
+ current_user = await crud .user .activate_totp (db = db , db_obj = current_user , totp_in = totp_in )
167
+ current_user = await crud .user .update_totp_counter (db = db , db_obj = current_user , new_counter = new_counter )
166
168
return {"msg" : "TOTP enabled. Do not lose your recovery code." }
167
169
168
170
169
171
@router .delete ("/totp" , response_model = schemas .Msg )
170
- def disable_totp_authentication (
172
+ async def disable_totp_authentication (
171
173
* ,
172
- db : Session = Depends (deps .get_db ),
174
+ db : AgnosticDatabase = Depends (deps .get_db ),
173
175
data_in : schemas .UserUpdate ,
174
176
current_user : models .User = Depends (deps .get_current_active_user ),
175
177
) -> Any :
176
178
"""
177
179
Disable TOTP.
178
180
"""
179
181
if current_user .hashed_password :
180
- user = crud .user .authenticate (db , email = current_user .email , password = data_in .original )
182
+ user = await crud .user .authenticate (db , email = current_user .email , password = data_in .original )
181
183
if not data_in .original or not user :
182
184
raise HTTPException (status_code = 400 , detail = "Unable to authenticate or deactivate TOTP." )
183
- crud .user .deactivate_totp (db = db , db_obj = current_user )
185
+ await crud .user .deactivate_totp (db = db , db_obj = current_user )
184
186
return {"msg" : "TOTP disabled." }
185
187
186
188
187
189
@router .post ("/refresh" , response_model = schemas .Token )
188
- def refresh_token (
189
- db : Session = Depends (deps .get_db ),
190
+ async def refresh_token (
191
+ db : AgnosticDatabase = Depends (deps .get_db ),
190
192
current_user : models .User = Depends (deps .get_refresh_user ),
191
193
) -> Any :
192
194
"""
193
195
Refresh tokens for future requests
194
196
"""
195
197
refresh_token = security .create_refresh_token (subject = current_user .id )
196
- crud .token .create (db = db , obj_in = refresh_token , user_obj = current_user )
198
+ await crud .token .create (db = db , obj_in = refresh_token , user_obj = current_user )
197
199
return {
198
200
"access_token" : security .create_access_token (subject = current_user .id ),
199
201
"refresh_token" : refresh_token ,
@@ -202,8 +204,8 @@ def refresh_token(
202
204
203
205
204
206
@router .post ("/revoke" , response_model = schemas .Msg )
205
- def revoke_token (
206
- db : Session = Depends (deps .get_db ),
207
+ async def revoke_token (
208
+ db : AgnosticDatabase = Depends (deps .get_db ),
207
209
current_user : models .User = Depends (deps .get_refresh_user ),
208
210
) -> Any :
209
211
"""
@@ -213,11 +215,11 @@ def revoke_token(
213
215
214
216
215
217
@router .post ("/recover/{email}" , response_model = Union [schemas .WebToken , schemas .Msg ])
216
- def recover_password (email : str , db : Session = Depends (deps .get_db )) -> Any :
218
+ async def recover_password (email : str , db : AgnosticDatabase = Depends (deps .get_db )) -> Any :
217
219
"""
218
220
Password Recovery
219
221
"""
220
- user = crud .user .get_by_email (db , email = email )
222
+ user = await crud .user .get_by_email (db , email = email )
221
223
if user and crud .user .is_active (user ):
222
224
tokens = security .create_magic_tokens (subject = user .id )
223
225
if settings .EMAILS_ENABLED :
@@ -227,9 +229,9 @@ def recover_password(email: str, db: Session = Depends(deps.get_db)) -> Any:
227
229
228
230
229
231
@router .post ("/reset" , response_model = schemas .Msg )
230
- def reset_password (
232
+ async def reset_password (
231
233
* ,
232
- db : Session = Depends (deps .get_db ),
234
+ db : AgnosticDatabase = Depends (deps .get_db ),
233
235
new_password : str = Body (...),
234
236
claim : str = Body (...),
235
237
magic_in : bool = Depends (deps .get_magic_token ),
@@ -239,7 +241,7 @@ def reset_password(
239
241
"""
240
242
claim_in = deps .get_magic_token (token = claim )
241
243
# Get the user
242
- user = crud .user .get (db , id = magic_in .sub )
244
+ user = await crud .user .get (db , id = magic_in .sub )
243
245
# Test the claims
244
246
if (
245
247
(claim_in .sub == magic_in .sub )
@@ -251,6 +253,5 @@ def reset_password(
251
253
# Update the password
252
254
hashed_password = security .get_password_hash (new_password )
253
255
user .hashed_password = hashed_password
254
- db .add (user )
255
- db .commit ()
256
+ await user .save ()
256
257
return {"msg" : "Password updated successfully." }
0 commit comments