11# Builtin library
2- from datetime import datetime
32import hashlib
3+ import aiohttp
44
55# Framework core library
66from starlette .requests import Request
77from fastapi import APIRouter , Header
88from fastapi .responses import JSONResponse
9+ import pyotp
910
1011# Local file
11- from util import random_content , json_body
12+ from util import json_body , token_tool
1213import util .pymongo_wrapper as DocumentDB
1314
1415router = APIRouter ()
1516
1617
17- @router .post ("/token/generate" , tags = ["V2" ])
18- async def v2_generate_auth_token (request : Request , cred : json_body .PasswordLoginBody ):
18+ @router .post ("/token/revoke" , tags = ["V2" ])
19+ async def v2_revoke_auth_token (request : Request , pa_token : str = Header (None )):
20+ mongo_client = DocumentDB .get_client ()
21+ db_client = mongo_client .get_database (DocumentDB .DB )
22+ token_deletion_query = DocumentDB .delete_one (
23+ collection = "TokenV3" ,
24+ find_filter = {"token_value" : pa_token },
25+ db_client = db_client )
26+ mongo_client .close ()
27+ if token_deletion_query is None :
28+ return JSONResponse (status_code = 500 , content = {"status" : "failed to remove the old token to database" , "pa_token" : pa_token })
29+ elif token_deletion_query .deleted_count == 0 :
30+ return JSONResponse (status_code = 404 , content = {"status" : "token not found" , "pa_token" : pa_token })
31+ elif token_deletion_query .deleted_count == 1 :
32+ return JSONResponse (status_code = 200 , content = {"status" : "deleted" , "pa_token" : pa_token })
33+
34+
35+ @router .post ("/password/verify" , tags = ["V2" ])
36+ async def v2_verify_auth_password (request : Request , cred : json_body .PasswordLoginBody ):
1937 mongo_client = DocumentDB .get_client ()
2038 db_client = mongo_client .get_database (DocumentDB .DB )
2139 credential_verify_query = DocumentDB .find_one (collection = "LoginV1" ,
@@ -30,51 +48,12 @@ async def v2_generate_auth_token(request: Request, cred: json_body.PasswordLogin
3048 content = {"status" : "not found or not match" ,
3149 "person_id" : cred .person_id ,
3250 "password" : cred .password })
33- while True :
34- # Checking if the same token already being use
35- # There is no do-while loop in Python
36- generated_token = random_content .generate_access_token ()
37- current_checking_query = DocumentDB .find_one (collection = "TokenV1" ,
38- find_filter = {"token_value" : generated_token },
39- db_client = db_client )
40- if current_checking_query is None :
41- break
42- create_at = int (datetime .now ().timestamp ())
43- expire_at = create_at + cred .token_lifespan
44- token_record_query = DocumentDB .insert_one (
45- collection = "TokenV3" ,
46- document_body = {
47- "structure_version" : 3 ,
48- "person_id" : cred .person_id ,
49- "token_value" : generated_token ,
50- "token_hash" : hashlib .sha512 (generated_token .encode ("utf-8" )).hexdigest (),
51- "creation_timestamp_int" : create_at ,
52- "expiration_timestamp_int" : expire_at
53- },
54- db_client = db_client )
55- if token_record_query is None :
56- return JSONResponse (status_code = 500 ,
57- content = {"status" : "token generated but failed to insert that token to database" })
51+ generated_token = token_tool .generate_pa_token_and_record (db_client = db_client ,
52+ person_id = cred .person_id ,
53+ token_lifespan = cred .token_lifespan )
5854 mongo_client .close ()
5955 return JSONResponse (status_code = 200 ,
60- content = {"status" : "success" , "pa_token" : generated_token , "expiration_timestamp" : expire_at })
61-
62-
63- @router .post ("/token/revoke" , tags = ["V2" ])
64- async def v2_revoke_auth_token (request : Request , pa_token : str = Header (None )):
65- mongo_client = DocumentDB .get_client ()
66- db_client = mongo_client .get_database (DocumentDB .DB )
67- token_deletion_query = DocumentDB .delete_one (
68- collection = "TokenV3" ,
69- find_filter = {"token_value" : pa_token },
70- db_client = db_client )
71- mongo_client .close ()
72- if token_deletion_query is None :
73- return JSONResponse (status_code = 500 , content = {"status" : "failed to remove the old token to database" , "pa_token" : pa_token })
74- elif token_deletion_query .deleted_count == 0 :
75- return JSONResponse (status_code = 404 , content = {"status" : "token not found" , "pa_token" : pa_token })
76- elif token_deletion_query .deleted_count == 1 :
77- return JSONResponse (status_code = 200 , content = {"status" : "deleted" , "pa_token" : pa_token })
56+ content = {"status" : "success" , "pa_token" : generated_token [0 ], "expiration_timestamp" : generated_token [1 ]})
7857
7958
8059# TODO: revoke existing session/token
@@ -112,3 +91,147 @@ async def v2_update_auth_password(request: Request, old_cred: json_body.Password
11291 "password" : old_cred .password })
11392 mongo_client .close ()
11493 return JSONResponse (status_code = 200 , content = {"status" : "success" , "voided" : old_cred .password })
94+
95+
96+ @router .post ("/totp/enable" , tags = ["V2" ])
97+ async def v2_enable_auth_totp (request : Request , cred : json_body .PasswordLoginBody ):
98+ mongo_client = DocumentDB .get_client ()
99+ db_client = mongo_client .get_database (DocumentDB .DB )
100+ # same as the traditional plain-password login
101+ credential_verify_query = DocumentDB .find_one (collection = "LoginV2" ,
102+ find_filter = {"person_id" : cred .person_id },
103+ db_client = db_client )
104+ print (credential_verify_query )
105+ if (credential_verify_query is None ) \
106+ or (hashlib .sha512 (cred .password .encode ("utf-8" )).hexdigest () != credential_verify_query ["password_hash" ]):
107+ mongo_client .close ()
108+ return JSONResponse (status_code = 403 ,
109+ content = {"status" : "user not found or not match" ,
110+ "person_id" : cred .person_id ,
111+ "password" : cred .password })
112+ if credential_verify_query ["totp_status" ] != "disabled" :
113+ mongo_client .close ()
114+ return JSONResponse (status_code = 200 ,
115+ content = {"status" : "Time-based OTP already enabled for this user" ,
116+ "person_id" : cred .person_id })
117+ new_secret_key = pyotp .random_base32 ()
118+ authenticator_url = pyotp .totp .TOTP (new_secret_key ).provisioning_uri (name = cred .person_id ,
119+ issuer_name = 'Plan-At' )
120+ credential_modify_query = DocumentDB .update_one (db_client = db_client ,
121+ collection = "LoginV2" ,
122+ find_filter = {"person_id" : cred .person_id },
123+ changes = {"$set" : {"totp_status" : "enabled" ,
124+ "totp_secret_key" : new_secret_key }})
125+ if credential_modify_query .matched_count != 1 and credential_modify_query .modified_count != 1 :
126+ return JSONResponse (status_code = 500 , content = {"status" : "failed to register the secret_key for totp in database" ,
127+ "matched_count" : credential_modify_query .matched_count ,
128+ "modified_count" : credential_modify_query .modified_count })
129+ mongo_client .close ()
130+ return JSONResponse (status_code = 200 ,
131+ content = {"status" : "Time-based OTP enabled for this user" ,
132+ "person_id" : cred .person_id ,
133+ "authenticator_url" : authenticator_url })
134+
135+
136+ @router .post ("/totp/disable" , tags = ["V2" ])
137+ async def v2_disable_auth_totp (request : Request , cred : json_body .PasswordLoginBody ):
138+ # Copy and Paste of /enable
139+ # Not require current output from Authenticator since the user might lose their
140+ mongo_client = DocumentDB .get_client ()
141+ db_client = mongo_client .get_database (DocumentDB .DB )
142+ # same as the traditional plain-password login
143+ credential_verify_query = DocumentDB .find_one (collection = "LoginV2" ,
144+ find_filter = {"person_id" : cred .person_id },
145+ db_client = db_client )
146+ print (credential_verify_query )
147+ if (credential_verify_query is None ) \
148+ or (hashlib .sha512 (cred .password .encode ("utf-8" )).hexdigest () != credential_verify_query ["password_hash" ]):
149+ mongo_client .close ()
150+ return JSONResponse (status_code = 403 ,
151+ content = {"status" : "user not found or not match" ,
152+ "person_id" : cred .person_id ,
153+ "password" : cred .password })
154+ # checking current totp status
155+ if credential_verify_query ["totp_status" ] != "enabled" :
156+ mongo_client .close ()
157+ return JSONResponse (status_code = 200 ,
158+ content = {"status" : "Time-based OTP not enabled for this user" ,
159+ "person_id" : cred .person_id })
160+ credential_modify_query = DocumentDB .update_one (db_client = db_client ,
161+ collection = "LoginV2" ,
162+ find_filter = {"person_id" : cred .person_id },
163+ changes = {"$set" : {"totp_status" : "disabled" ,
164+ "totp_secret_key" : "" }})
165+ if credential_modify_query .matched_count != 1 and credential_modify_query .modified_count != 1 :
166+ return JSONResponse (status_code = 500 , content = {"status" : "failed to delete existing secret_key for totp in database" ,
167+ "matched_count" : credential_modify_query .matched_count ,
168+ "modified_count" : credential_modify_query .modified_count })
169+ mongo_client .close ()
170+ return JSONResponse (status_code = 200 ,
171+ content = {"status" : "Time-based OTP disabled for this user" ,
172+ "person_id" : cred .person_id })
173+
174+
175+ @router .post ("/totp/verify" , tags = ["V2" ])
176+ async def v2_verify_auth_totp (request : Request , person_id : str , totp_code : str ):
177+ # Copy and Paste of /disable
178+ # Not require current output from Authenticator since the user might lose their
179+ mongo_client = DocumentDB .get_client ()
180+ db_client = mongo_client .get_database (DocumentDB .DB )
181+ # if len(str(int(totp_code))) != 6: # also verify if its actual int but not working if start with zero
182+ if len (totp_code ) != 6 : # also verify if its actual int but not working if start with zero
183+ return JSONResponse (status_code = 400 ,
184+ content = {"status" : "totp_code malformed" ,
185+ "totp_code" : totp_code })
186+ # same as the traditional plain-password login
187+ credential_verify_query = DocumentDB .find_one (collection = "LoginV2" ,
188+ find_filter = {"person_id" : person_id },
189+ db_client = db_client )
190+ print (credential_verify_query )
191+ if credential_verify_query is None :
192+ mongo_client .close ()
193+ return JSONResponse (status_code = 403 ,
194+ content = {"status" : "user not found or totp_code not match" ,
195+ "person_id" : person_id ,
196+ "totp_code" : totp_code })
197+ # Checking current totp status
198+ if credential_verify_query ["totp_status" ] != "enabled" :
199+ mongo_client .close ()
200+ return JSONResponse (status_code = 200 ,
201+ content = {"status" : "Time-based OTP not enabled for this user" ,
202+ "person_id" : person_id })
203+
204+ if not pyotp .TOTP (credential_verify_query ["totp_secret_key" ]).verify (totp_code ):
205+ mongo_client .close ()
206+ return JSONResponse (status_code = 403 ,
207+ content = {"status" : "user not found or totp_code not match" ,
208+ "person_id" : person_id ,
209+ "totp_code" : totp_code })
210+ generated_token = token_tool .generate_pa_token_and_record (db_client = db_client ,
211+ person_id = person_id ,
212+ token_lifespan = (60 * 60 * 24 * 1 ))
213+ mongo_client .close ()
214+ return JSONResponse (status_code = 200 ,
215+ content = {"status" : "success" ,
216+ "person_id" : person_id ,
217+ "pa_token" : generated_token [0 ],
218+ "expiration_timestamp" : generated_token [1 ]})
219+
220+
221+ @router .post ("/github/enable" , tags = ["V2" ])
222+ async def v2_enable_auth_github (request : Request , req_body : json_body .GitHubOAuthCode , pa_token : str = Header (None )):
223+ github_session = aiohttp .ClientSession ()
224+ a = await github_session .post (f"https://github.com/login/oauth/access_token?client_id={ 1 } &client_secret={ 2 } &code={ 3 } " )
225+ print (a .status , a .text ())
226+ a = a .json ()
227+ return JSONResponse (status_code = 200 , content = {"status" : "success" , "code" : req_body .code })
228+
229+
230+ @router .post ("/github/disable" , tags = ["V2" ])
231+ async def v2_disable_auth_github (request : Request , req_body : json_body .GitHubOAuthCode , pa_token : str = Header (None )):
232+ pass
233+
234+
235+ @router .post ("/github/verify" , tags = ["V2" ])
236+ async def v2_verify_auth_github (request : Request , req_body : json_body .GitHubOAuthCode ):
237+ pass
0 commit comments