Browse Source

refactor API

boyska 2 years ago
parent
commit
fff408bd24
1 changed files with 35 additions and 29 deletions
  1. 35 29
      tresetter.py

+ 35 - 29
tresetter.py

@@ -9,7 +9,7 @@ from typing import Optional
 
 import redis
 from fastapi import FastAPI, APIRouter, HTTPException, Cookie, Request
-from fastapi.responses import JSONResponse, RedirectResponse
+from fastapi.responses import Response, JSONResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, BaseSettings
 
@@ -92,15 +92,18 @@ def delete_session(session_id: str):
 # Models {{{
 
 
-class LoginData(BaseModel):
+class UserData(BaseModel):
     username: str
-    password: str
 
 
 class ChangeData(BaseModel):
     password: str
 
 
+class LoginData(UserData, ChangeData):
+    pass
+
+
 class SuccessData(BaseModel):
     success: bool = True
 
@@ -127,11 +130,28 @@ def password_generate():
     return "".join(random.choices(symbols, k=10))
 
 
+def change_password(username: str, new_password: str) -> bool:
+    p = Popen(
+        [settings.change_password_exe],
+        env={"CHANGE_USERNAME": username, "CHANGE_PASSWORD": new_password},
+    )
+
+    try:
+        p.communicate()
+    except CalledProcessError:
+        return False
+
+    if p.returncode != 0:
+        return False
+    return True
+
+
 # end external commands }}}
 
 
 @router.get("/")
 async def home(request: Request, session_id: str = Cookie(None)):
+    """redirects to the user to the home"""
     try:
         get_session(session_id)
     except SessionNotFoundException:
@@ -141,6 +161,9 @@ async def home(request: Request, session_id: str = Cookie(None)):
 
 @router.post("/login", tags=["session"])
 async def login(req: LoginData):
+    """
+    performs login
+    """
 
     ok = validate(req.username, req.password)
     if not ok:
@@ -150,27 +173,23 @@ async def login(req: LoginData):
             "username": req.username,
         }
     )
-    response = JSONResponse(
-        content={
-            "status": "ok",
-            "username": req.username,
-        }
-    )
+    response = Response()
     response.set_cookie(key="session_id", value=session_id)
     return response
 
 
 @router.get("/whoami", tags=["session"])
 async def whoami(session_id: str = Cookie(None)):
+    """Confirm login information"""
     session = get_session(session_id)
-    return JSONResponse(content={"username": session["username"]})
+    return UserData(username=session["username"])
 
 
 @router.post("/logout", tags=["session"])
-async def logout(session_id: str = Cookie(None)):
+async def logout(session_id: str = Cookie(None)) -> BaseModel:
     get_session(session_id)
     delete_session(session_id)
-    return "OKI"
+    return BaseModel()
 
 
 @router.post("/generate", tags=["password"])
@@ -179,32 +198,19 @@ async def generate(session_id: str = Cookie(None)):
     session["proposed_password"] = password_generate()
     set_session(session_id, session)
 
-    return JSONResponse(content={"password": session["proposed_password"]})
+    return ChangeData(password=session["proposed_password"])
 
 
 @router.post("/change", tags=["password"])
-async def change(req: ChangeData, session_id: str = Cookie(None)):
+async def change(req: ChangeData, session_id: str = Cookie(None)) -> SuccessData:
     session = get_session(session_id)
     if "proposed_password" not in session:
         raise HTTPException(status_code=400, detail="You must generate it first")
     if req.password != session["proposed_password"]:
         raise HTTPException(status_code=409)
 
-    p = Popen(
-        [settings.change_password_exe],
-        env={"CHANGE_USERNAME": session["username"], "CHANGE_PASSWORD": req.password},
-    )
-
-    try:
-        p.communicate()
-    except CalledProcessError:
-        fail = True
-    else:
-        fail = p.returncode != 0
-    if fail:
-        return SuccessData(success=False)
-
-    return SuccessData()
+    success = change_password(session["username"], req.password)
+    return SuccessData(success=success)
 
 
 app.include_router(router)