diff --git a/tresetter.py b/tresetter.py index 2958bae..3092c5b 100644 --- a/tresetter.py +++ b/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)