From e070a96fb6dc88a8acbac9cd0a9ce8e714f5cb37 Mon Sep 17 00:00:00 2001
From: boyska <boyska@riseup.net>
Date: Fri, 20 May 2022 19:07:02 +0200
Subject: [PATCH] use a kdf

---
 tresetter.py | 37 ++++++++++++++++++++++++++++++++-----
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/tresetter.py b/tresetter.py
index d530d17..a728132 100644
--- a/tresetter.py
+++ b/tresetter.py
@@ -6,10 +6,12 @@ import json
 import uuid
 from subprocess import Popen, CalledProcessError, check_output
 from typing import Optional
+import hashlib
+import base64
 
 import redis
 from fastapi import FastAPI, APIRouter, HTTPException, Cookie, Request
-from fastapi.responses import Response, JSONResponse, RedirectResponse
+from fastapi.responses import Response, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, BaseSettings
 
@@ -201,22 +203,47 @@ async def logout(session_id: str = Cookie(None)) -> BaseModel:
     delete_session(session_id)
     return BaseModel()
 
+KDF_SALT_SIZE = 16
+
+
+def kdf_gen(password, salt=None) -> str:
+    if salt is None:
+        salt = random.randbytes(KDF_SALT_SIZE)
+    if hasattr(password, 'encode'):
+        password = password.encode('utf8')
+    raw = hashlib.scrypt(password, n=2, r=1, p=1, salt=salt)
+    with_salt = salt + raw
+    return base64.b64encode(with_salt).decode('ascii')
+
+def kdf_get_salt(hashed: str):
+    hashed_str = hashed.decode('ascii') if hasattr(hashed, 'decode') else hashed
+    with_salt = base64.b64decode(hashed_str)
+    salt = with_salt[:KDF_SALT_SIZE]
+    return salt
+
+def kdf_verify(hashed: str, password: str) -> bool:
+    salt = kdf_get_salt(hashed)
+    hashed2 = kdf_gen(password, salt=salt)
+    return hashed == hashed2
+
 
 @router.post("/generate", tags=["password"])
 async def generate(session_id: str = Cookie(None)):
     session = get_session(session_id)
-    session["proposed_password"] = password_generate()
+    proposed_password = password_generate()
+    session["proposed_password_hash"] = kdf_gen(proposed_password)
     set_session(session_id, session)
 
-    return ChangeData(password=session["proposed_password"])
+    return ChangeData(password=proposed_password)
 
 
 @router.post("/change", tags=["password"])
 async def change(req: ChangeData, session_id: str = Cookie(None)) -> SuccessData:
     session = get_session(session_id)
-    if "proposed_password" not in session:
+    if "proposed_password_hash" not in session:
         raise HTTPException(status_code=400, detail="You must generate it first")
-    if req.password != session["proposed_password"]:
+    hashed = session["proposed_password_hash"]
+    if not kdf_verify(hashed, req.password):
         raise HTTPException(status_code=409)
 
     success = change_password(session["username"], req.password)