From aa6e102b073bd045c63bcffe951b03e99bfa0917 Mon Sep 17 00:00:00 2001
From: Balakrishnan Balasubramanian <git.user@balki.me>
Date: Tue, 6 Jun 2023 00:00:36 -0400
Subject: [PATCH] CLI for generating password hashes and add unittest for
 pwhash

---
 mail4one/pwhash.py      | 13 +++++++++++-
 mail4one/pwhash_test.py | 44 +++++++++++++++++++++++++++--------------
 2 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/mail4one/pwhash.py b/mail4one/pwhash.py
index 5d3ccdc..24cc6af 100644
--- a/mail4one/pwhash.py
+++ b/mail4one/pwhash.py
@@ -29,7 +29,7 @@ def gen_pwhash(password: str) -> str:
 
 class PWInfo:
 
-    def __init__(self, salt, sh):
+    def __init__(self, salt: bytes, sh: bytes):
         self.salt = salt
         self.scrypt_hash = sh
 
@@ -52,3 +52,14 @@ def check_pass(password: str, pwinfo: PWInfo) -> bool:
                                         n=SCRYPT_N,
                                         r=SCRYPT_R,
                                         p=SCRYPT_P)
+
+
+if __name__ == '__main__':
+    import sys
+    if len(sys.argv) == 2:
+        print(gen_pwhash(sys.argv[1]))
+    elif len(sys.argv) == 3:
+        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)
diff --git a/mail4one/pwhash_test.py b/mail4one/pwhash_test.py
index 778b69f..adb76d6 100644
--- a/mail4one/pwhash_test.py
+++ b/mail4one/pwhash_test.py
@@ -1,22 +1,36 @@
-from .pwhash import *
+from .pwhash import gen_pwhash, parse_hash, check_pass, SALT_LEN
+import unittest
 
-def check_pass_from_hash(password: str, pwhash: str) -> bool:
-    try:
+
+class TestPWHash(unittest.TestCase):
+
+    def test_expected_usage(self):
+        password = "Blah Blah ABCD"
+        pwhash = gen_pwhash(password)
         pwinfo = parse_hash(pwhash)
-    except:
-        return False
-    return check_pass(password, pwinfo)
+        self.assertEqual(len(pwinfo.salt), SALT_LEN)
+        self.assertEqual(len(pwinfo.scrypt_hash), 64)
+        self.assertTrue(check_pass(password, pwinfo),
+                        "check pass with correct password")
+        self.assertFalse(check_pass("foobar", pwinfo),
+                         "check pass with wrong password")
 
-test_hash = "AFWMBONQ2XGHWBTKVECDBBJWYEMS4DFIXIJML4VP76JQT5VWVLALE3KVKFEBAGWG3DOY53DK3H2EACWOBHJFYAIHDA3OFDQN2UAXI5TLBFOW4O2GWXNBGQ5QFMOJ5Z27HGYNO73DS5WPX2INNE47EGI6Z5UAAQAAAAEAAAIA"
+    def test_hardcoded_hash(self):
+        test_hash = "".join((l.strip() for l in """
+        AFWMBONQ2XGHWBTKVECDBBJWYEMS4DFIXIJML4VP76JQT5VWVLALE3KV
+        KFEBAGWG3DOY53DK3H2EACWOBHJFYAIHDA3OFDQN2UAXI5TLBFOW4O2G
+        WXNBGQ5QFMOJ5Z27HGYNO73DS5WPX2INNE47EGI6Z5UAAQAAAAEAAAIA
+        """.splitlines()))
+        pwinfo = parse_hash(test_hash)
+        self.assertTrue(check_pass("helloworld", pwinfo),
+                        "check pass with correct password")
+        self.assertFalse(check_pass("foobar", pwinfo),
+                         "check pass with wrong password")
 
-def main():
-    print(gen_pwhash("helloworld"))
-    print("------------")
-    print(check_pass_from_hash("hElloworld", test_hash))
-    print(check_pass_from_hash("helloworld", "foobar"))
-    print("------------")
-    print(check_pass_from_hash("helloworld", test_hash))
+    def test_invalid_hash(self):
+        with self.assertRaises(Exception):
+            parse_hash("sdlfkjdsklfjdsk")
 
 
 if __name__ == '__main__':
-    main()
+    unittest.main()