Introduce types.

1 files changed, 36 insertions, 23 deletions

diff --git a/password.py b/password.py
index 5e1f37e..16192be 100755
--- a/password.py
+++ b/password.py
@@ -3,59 +3,73 @@ import hashlib
 import json
 import os
 import random
+from typing import (
+    TypedDict,
+)
 
-def gen_salt(length=10):
+
+def gen_salt(length: int = 10) -> str:
     # TODO is this a sufficient source of randomness
     return bytearray(random.randint(0, 256) for _ in range(length)).hex()
 
+
 # Manual list of entries, to stop someone from executing arbitrary
 # code by modyfying password database
 hash_methods = {
-        'sha256': hashlib.sha256
-        }
+    'sha256': hashlib.sha256
+}
+
+
+class PasswordEntry(TypedDict):
+    hash: str
+    salt: str
+    # One of the keys of hash_methods
+    method: str
+
 
 class Passwords:
-    def __init__(self, fname):
+    def __init__(self, fname: os.PathLike):
         self.fname = fname
+        self.db: dict[str, PasswordEntry]
         try:
             with open(fname) as f:
                 self.db = json.load(f)
-        except:
+        except Exception:
             self.db = {}
 
-    def save(self):
+    def save(self) -> None:
         try:
-            with open(self.fname + '.tmp', 'w') as f:
+            with open(os.fspath(self.fname) + '.tmp', 'w') as f:
                 json.dump(self.db, f)
                 f.write('\n')
-            os.rename(self.fname + '.tmp', self.fname)
-        except e:
-            print('Saving password failed {e}')
+            os.rename(os.fspath(self.fname) + '.tmp', self.fname)
+        except Exception as e:
+            print(f'Saving password failed {e}')
 
-    def add(self, username, password):
+    def add(self, username: str, password: str) -> None:
         if cur := self.db.get(username):
             salt = cur['salt']
             hashed = hashlib.sha256((salt + password).encode('UTF-8'))
             self.db[username] = {
-                    'hash': hashed.hexdigest(),
-                    'salt': salt,
-                    'method': 'sha256',
-                    }
+                'hash': hashed.hexdigest(),
+                'salt': salt,
+                'method': 'sha256',
+            }
         else:
             salt = gen_salt()
             hashed = hashlib.sha256((salt + password).encode('UTF-8'))
             self.db[username] = {
-                    'hash': hashed.hexdigest(),
-                    'salt': salt,
-                    'method': 'sha256'
-                    }
+                'hash': hashed.hexdigest(),
+                'salt': salt,
+                'method': 'sha256'
+            }
 
-    def validate(self, username, password):
+    def validate(self, username: str, password: str) -> bool:
         # These shall fail when key is missing
         data = self.db[username]
         proc = hash_methods[data['method']]
-        return data['hash'] == proc((data['salt'] + password).encode('UTF-8')).hexdigest()
-
+        digest = proc((data['salt'] + password).encode('UTF-8')).hexdigest()
+        return data['hash'] == digest
 
 
 if __name__ == '__main__':
@@ -77,7 +91,6 @@ if __name__ == '__main__':
 
     args = parser.parse_args()
 
-
     passwords = Passwords(args.file)
     if args.cmd == 'add':
         passwords.add(args.username, args.password)