From 9b0400583c1f13db4fd8e5db9c906dc631b2b1c7 Mon Sep 17 00:00:00 2001
From: Balakrishnan Balasubramanian <git.user@balki.me>
Date: Wed, 7 Jun 2023 16:24:06 -0400
Subject: [PATCH] pwhash: Don't encode script parameters in hash

---
 mail4one/pwhash.py      | 38 ++++++++++++++++++++++----------------
 mail4one/pwhash_test.py | 10 +++++-----
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/mail4one/pwhash.py b/mail4one/pwhash.py
index 24cc6af..09396be 100644
--- a/mail4one/pwhash.py
+++ b/mail4one/pwhash.py
@@ -1,18 +1,20 @@
 import os
 from hashlib import scrypt
-from struct import pack, unpack
 from base64 import b32encode, b32decode
 
 # Links
 # https://pkg.go.dev/golang.org/x/crypto/scrypt#Key
 # https://crypto.stackexchange.com/a/35434
 
+# Doubling N causes Memory limit exceeded
 SCRYPT_N = 16384
 SCRYPT_R = 8
 SCRYPT_P = 1
+
+# If any of above parameters change, version will be incremented
 VERSION = b'\x01'
-SALT_LEN = 32
-PACK_FMT = f"<c{SALT_LEN}s64slhh"
+SALT_LEN = 30
+KEY_LEN = 64  # This is python default
 
 
 def gen_pwhash(password: str) -> str:
@@ -21,10 +23,9 @@ def gen_pwhash(password: str) -> str:
                 salt=salt,
                 n=SCRYPT_N,
                 r=SCRYPT_R,
-                p=SCRYPT_P)
-    pack_bytes = pack(PACK_FMT, VERSION, salt, sh, SCRYPT_N, SCRYPT_R,
-                      SCRYPT_P)
-    return b32encode(pack_bytes).decode()
+                p=SCRYPT_P,
+                dklen=KEY_LEN)
+    return b32encode(VERSION + salt + sh).decode()
 
 
 class PWInfo:
@@ -34,14 +35,17 @@ class PWInfo:
         self.scrypt_hash = sh
 
 
-def parse_hash(pwhash: str) -> PWInfo:
-    decoded = b32decode(pwhash.encode())
-    ver, salt, sh, n, r, p = unpack(PACK_FMT, decoded)
-    if not (ver, n, r, p, len(salt)) == (VERSION, SCRYPT_N, SCRYPT_R, SCRYPT_P,
-                                         SALT_LEN):
+def parse_hash(pwhash_str: str) -> PWInfo:
+    pwhash = b32decode(pwhash_str.encode())
+
+    if not len(pwhash) == 1 + SALT_LEN + KEY_LEN:
         raise Exception(
-            f"Invalid hash: {ver=}, {n=}, {r=}, {p=}, f{len(salt)=} != {VERSION=}, {SCRYPT_N=}, {SCRYPT_R=}, {SCRYPT_P=}, {SALT_LEN=}"
-        )
+            f"Invalid hash size, {len(pwhash)} !=  {1 + SALT_LEN + KEY_LEN}")
+
+    if (ver := pwhash[0:1]) != VERSION:
+        raise Exception(f"Invalid hash version, {ver!r} !=  {VERSION!r}")
+
+    salt, sh = pwhash[1:SALT_LEN + 1], pwhash[-KEY_LEN:]
     return PWInfo(salt, sh)
 
 
@@ -51,7 +55,8 @@ def check_pass(password: str, pwinfo: PWInfo) -> bool:
                                         salt=pwinfo.salt,
                                         n=SCRYPT_N,
                                         r=SCRYPT_R,
-                                        p=SCRYPT_P)
+                                        p=SCRYPT_P,
+                                        dklen=KEY_LEN)
 
 
 if __name__ == '__main__':
@@ -62,4 +67,5 @@ if __name__ == '__main__':
         ok = check_pass(sys.argv[1], parse_hash(sys.argv[2]))
         print("OK" if ok else "NOT OK")
     else:
-        print("Usage: python3 -m mail4one.pwhash <password> [password_hash]", file=sys.stderr)
+        print("Usage: python3 -m mail4one.pwhash <password> [password_hash]",
+              file=sys.stderr)
diff --git a/mail4one/pwhash_test.py b/mail4one/pwhash_test.py
index adb76d6..b18b96b 100644
--- a/mail4one/pwhash_test.py
+++ b/mail4one/pwhash_test.py
@@ -1,4 +1,4 @@
-from .pwhash import gen_pwhash, parse_hash, check_pass, SALT_LEN
+from .pwhash import gen_pwhash, parse_hash, check_pass, SALT_LEN, KEY_LEN
 import unittest
 
 
@@ -9,7 +9,7 @@ class TestPWHash(unittest.TestCase):
         pwhash = gen_pwhash(password)
         pwinfo = parse_hash(pwhash)
         self.assertEqual(len(pwinfo.salt), SALT_LEN)
-        self.assertEqual(len(pwinfo.scrypt_hash), 64)
+        self.assertEqual(len(pwinfo.scrypt_hash), KEY_LEN)
         self.assertTrue(check_pass(password, pwinfo),
                         "check pass with correct password")
         self.assertFalse(check_pass("foobar", pwinfo),
@@ -17,9 +17,9 @@ class TestPWHash(unittest.TestCase):
 
     def test_hardcoded_hash(self):
         test_hash = "".join((l.strip() for l in """
-        AFWMBONQ2XGHWBTKVECDBBJWYEMS4DFIXIJML4VP76JQT5VWVLALE3KV
-        KFEBAGWG3DOY53DK3H2EACWOBHJFYAIHDA3OFDQN2UAXI5TLBFOW4O2G
-        WXNBGQ5QFMOJ5Z27HGYNO73DS5WPX2INNE47EGI6Z5UAAQAAAAEAAAIA
+        AFTY5EVN7AX47ZL7UMH3BETYWFBTAV3XHR73CEFAJBPN2NIHPWD
+        ZHV2UQSMSPHSQQ2A2BFQBNC77VL7F2UKATQNJZGYLCSU6C43UQD
+        AQXWXSWNGAEPGIMG2F3QDKBXL3MRHY6K2BPID64ZR6LABLPVSF
         """.splitlines()))
         pwinfo = parse_hash(test_hash)
         self.assertTrue(check_pass("helloworld", pwinfo),