-
Notifications
You must be signed in to change notification settings - Fork 4
/
main.py
67 lines (47 loc) · 2.06 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
from pydantic import BaseModel
from time import time
from fastapi import FastAPI, HTTPException
from totp import generate_secret, check_code, time_window
from users import users
class User(BaseModel):
user_id: str
totp_secret: str
misses_since_success: int = 0
class TotpCheckResult(BaseModel):
success: bool
misses_since_success: int
app = FastAPI()
@app.get("/v1/")
def read_root():
return {"description": "A Simple FastAPI TOTP Server v1"}
@app.post("/v1/users/create/{user_id}", response_model=User, status_code=201)
def create_user(user_id: str):
"""Try to create new user; replies with new user object along with a secret, or appropriate error"""
if user_id in users:
raise HTTPException(status_code=409, detail="User '{}' already exists".format(user_id))
user = User(user_id=user_id, totp_secret=generate_secret())
users[user_id] = user
return user
@app.put("/v1/users/check_totp/{user_id}/{code}", response_model=TotpCheckResult)
def check_totp(user_id: str, code: str):
"""Checks if the code of a user is correct for the current moment.
Arguments:
user_id -- user to check against
code -- 4-digit string TOTP code.
Returns True if the code is valid for the current or previous time window
"""
if not users.get(user_id):
raise HTTPException(status_code=404, detail="User '{}' not found".format(user_id))
check_response = TotpCheckResult(success=False, misses_since_success=1000)
try:
# success if code is valid for current or immediately preceding time window
check_response.success = check_code(users[user_id].totp_secret, code) or \
check_code(users[user_id].totp_secret, code, int(time()) - time_window)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
if not check_response.success:
users[user_id].misses_since_success += 1
else:
users[user_id].misses_since_success = 0
check_response.misses_since_success = users[user_id].misses_since_success
return check_response