import hashlib
import secrets
import os, sys
import getpass

HASHTYPE = "sha256"
HASHITERS = 1000

SECTS = (64, 32, 32)
SECTLEN = sum(SECTS)
#assert sum(SECTS) == SECTLEN, "Invalid division of sections." # implied in defenition

file = os.path.expanduser("~/.pswd.db")

def write(rows=[]):
    with open(file, "wb") as f:
        for row in rows:
            for col in row:
                f.write(col)

def add(uname, pswd):
    #data = hashlib.md5(uname.encode("utf-8")).digest() # 16 byte MD5 uname hash
    data = uname.encode("utf-8")[:SECTS[0]].ljust(SECTS[0], b"\x00") # 32 byte UTF8 name, null-padded

    salt = secrets.token_bytes(SECTS[2])
    pswdhash = hashlib.pbkdf2_hmac(HASHTYPE, pswd.encode("utf-8"), salt, HASHITERS)
    data += pswdhash + salt

    assert len(data) == SECTLEN, "Invalid data section length."

    with open(file, mode="ab") as f:
        f.write(data)

def read():
    with open(file, mode="rb") as f:
        raw = f.read()
    if len(raw) % SECTLEN != 0:
        print("WARNING: password file has invalid length; discarding end", file=sys.stderr)
    rows = []
    for x in range(0, len(raw), SECTLEN):
        row = []
        c = 0
        for s in SECTS:
            row.append(raw[x+c:x+c+s])
            c += s
        rows.append(row)
    write(rows)
    return rows

fails = 0
success = False
while fails < 5 and not success:
    user = input("Username: ")

    #if hashlib.md5(user.encode("utf-8")).digest() in map(lambda x: x[0], read()):
    if user.encode("utf-8")[:SECTS[0]].ljust(SECTS[0], b"\x00") in map(lambda x: x[0], read()):
        # user is in database
        #index = list(map(lambda x: x[0], read())).index(hashlib.md5(user.encode("utf-8")).digest())
        index = list(map(lambda x: x[0], read())).index(user.encode("utf-8")[:SECTS[0]].ljust(SECTS[0], b"\x00"))
        userhash, pswdhash, salt = read()[index]

        pswd = getpass.getpass("Password: ")
        if hashlib.pbkdf2_hmac(HASHTYPE, pswd.encode("utf-8"), salt, HASHITERS) == pswdhash:
            success = True
        else:
            print("Incorrect password.")
            fails += 1

    else:
        yn = input("Unrecognised username. Would you like to create an account? (y/n) ")
        if yn.lower()[0] == "y":
            pswd1 = getpass.getpass("Password: ")
            pswd2 = getpass.getpass("Confirm: ")
            if pswd1 != pswd2:
                print("Non-matching passwords. Aborting.")
            else:
                add(user, pswd1)
        else:
            print("Aborting.")

if success:
    print(f"Welcome, {user}!")
    # more program here
else:
    print("Locked out.")