added a few tests and changed add_users

auth.py

@@ -0,0 +1,80 @@
+import getpass
+import pathlib
+import pickle
+import random
+import string
+import tempfile
+
+PWDB_FLNAME = pathlib.Path('pwdb.pkl')
+CHARS = string.ascii_letters + string.digits + string.punctuation
+class PasswordError(IOError): pass
+class UsernameError(IOError): pass
+
+def get_credentials():
+    username = input('Enter your username: ')
+    password = getpass.getpass('Enter your password: ')
+    return (username, password)
+
+def authenticate(username, pass_text, pwdb):
+    if username in pwdb:
+        salt = pwdb[username][1]
+        if pwhash(pass_text, salt) == pwdb[username][0]:
+            return True
+    return False
+
+def add_user(username, password, salt, pwdb, pwdb_file):
+    if username in pwdb:
+        raise Exception('Username already exists [%s]' %username)
+    elif not username:
+        raise UsernameError('Please type in some username!')
+    elif not password:
+        raise PasswordError('Please type in some password!')
+    else:
+        pwdb[username] = (pwhash(password,salt), salt)
+        write_pwdb(pwdb, pwdb_file)
+
+def read_pwdb(pwdb_file):
+    try:
+        pwdb = pickle.load(pwdb_file)
+        pwdb_file.seek(0)
+    except EOFError:
+        pwdb = {}
+    return pwdb
+
+def write_pwdb(pwdb, pwdb_file):
+    pickle.dump(pwdb, pwdb_file)
+
+def pwhash(pass_text, salt):
+    hash_ = 0
+    full_pass_text = pass_text + salt
+    for idx, char in enumerate(full_pass_text):
+        # use idx as a multiplier, so that shuffling the characters returns a
+        # different hash
+        hash_ += (idx+1)*ord(char)
+    return hash_
+
+def get_salt():
+    salt_chars = random.choices(CHARS, k=10)
+    return ''.join(salt_chars)
+
+if __name__ == '__main__':
+    pwdb_path = tempfile.gettempdir() / PWDB_FLNAME
+    try:
+        pwdb_file = open(pwdb_path, 'rb+')
+    except FileNotFoundError:
+        pwdb_file = open(pwdb_path, 'wb+')
+
+    username, password = get_credentials()
+    pwdb = read_pwdb(pwdb_file)
+
+    if authenticate(username, password, pwdb):
+        print('Authentication succeeded!')
+        print(pwdb)
+    else:
+        print('Wrong username or password')
+        ans = input('Create new user [y/n]? ')
+        if ans == 'y':
+            salt = get_salt()
+            add_user(username, password, salt, pwdb, pwdb_file)
+        else:
+            print('Exit!')

test_auth.py

@@ -0,0 +1,116 @@
+import auth
+import getpass
+import pathlib
+import pickle
+import random
+import string
+import tempfile
+
+PWDB_FLNAME = pathlib.Path('test_pwdb.pkl')
+
+def test_right_name_right_password():
+    salt = auth.get_salt()
+    password = 'real_password'
+    pwdb = {'real_name':  (auth.pwhash(password,salt), salt)}
+    username = 'real_name'
+    pass_text = 'real_password'
+    assert auth.authenticate(username, pass_text, pwdb)
+
+def test_right_name_wrong_password():
+    salt = auth.get_salt()
+    password = 'real_password'
+    pwdb = {'real_name':  (auth.pwhash(password,salt), salt)}
+    username = 'real_name'
+    pass_text = 'wrong_password'
+    assert not auth.authenticate(username, pass_text, pwdb)
+
+def test_wrong_name_right_password():
+    salt = auth.get_salt()
+    password = 'real_password'
+    pwdb = {'real_name':  (auth.pwhash(password,salt), salt)}
+    username = 'wrong_name'
+    pass_text = 'real_password'
+    assert not auth.authenticate(username, pass_text, pwdb)
+
+def test_wrong_name_wrong_password():
+    salt = auth.get_salt()
+    password = 'real_password'
+    pwdb = {'real_name':  (auth.pwhash(password,salt), salt)}
+    username = 'wrong_name'
+    pass_text = 'wrong_password'
+    assert not auth.authenticate(username, pass_text, pwdb)
+
+def test_empy_database():
+    salt = auth.get_salt()
+    pwdb = {}
+    username = 'some_name'
+    pass_text = 'some_password'
+    assert not auth.authenticate(username, pass_text, pwdb)
+
+
+def test_user_already_exists():
+    username = 'old_name'
+    password = 'old_password'
+    salt = auth.get_salt()
+    pwdb = {'old_name':  (auth.pwhash('old_password',salt), salt)}
+    pwdb_path = tempfile.gettempdir() / PWDB_FLNAME
+    pwdb_file = open(pwdb_path, 'wb')
+    pickle.dump(pwdb, pwdb_file)
+    salt = auth.get_salt()
+    try:
+        auth.add_user(username, password, salt, pwdb, pwdb_file)
+        assert False
+    except Exception as _:
+        assert True
+
+
+def test_user_not_exists():
+    username = 'new_name'
+    password = 'new_password'
+    salt = auth.get_salt()
+    pwdb = {'old_name':  (auth.pwhash('old_password',salt), salt)}
+    pwdb_path = tempfile.gettempdir() / PWDB_FLNAME
+    with open(pwdb_path, 'wb+') as pwdb_file:
+        pickle.dump(pwdb, pwdb_file)
+    salt = auth.get_salt()
+    try:
+        with open(pwdb_path, 'wb+') as pwdb_file:
+            auth.add_user(username, password, salt, pwdb, pwdb_file)
+        with open(pwdb_path, 'rb+') as pwdb_file:
+            pwdb = pickle.load(pwdb_file)
+        print(pwdb)
+        assert pwdb[username] == (auth.pwhash(password,salt), salt)
+    except:
+        assert False
+
+def test_empty_username():
+    username = ''
+    password = 'new_password'
+    salt = auth.get_salt()
+    pwdb = {'old_name':  (auth.pwhash('old_password',salt), salt)}
+    pwdb_path = tempfile.gettempdir() / PWDB_FLNAME
+    with open(pwdb_path, 'wb+') as pwdb_file:
+        pickle.dump(pwdb, pwdb_file)
+    salt = auth.get_salt()
+    try:
+        with open(pwdb_path, 'wb+') as pwdb_file:
+            auth.add_user(username, password, salt, pwdb, pwdb_file)
+        assert False
+    except auth.UsernameError as _:
+        assert True
+
+def test_empty_password():
+    username = 'new_username'
+    password = ''
+    salt = auth.get_salt()
+    pwdb = {'old_name':  (auth.pwhash('old_password',salt), salt)}
+    pwdb_path = tempfile.gettempdir() / PWDB_FLNAME
+    with open(pwdb_path, 'wb+') as pwdb_file:
+        pickle.dump(pwdb, pwdb_file)
+    salt = auth.get_salt()
+    try:
+        with open(pwdb_path, 'wb+') as pwdb_file:
+            auth.add_user(username, password, salt, pwdb, pwdb_file)
+        assert False
+    except auth.PasswordError as _:
+        assert True