#!/usr/bin/env python3
import hashlib
import json
import os
import random

def gen_salt(length=10):
    # 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
        }

class Passwords:
    def __init__(self, fname):
        self.fname = fname
        try:
            with open(fname) as f:
                self.db = json.load(f)
        except:
            self.db = {}

    def save(self):
        try:
            with open(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}')

    def add(self, username, password):
        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',
                    }
        else:
            salt = gen_salt()
            hashed = hashlib.sha256((salt + password).encode('UTF-8'))
            self.db[username] = {
                    'hash': hashed.hexdigest(),
                    'salt': salt,
                    'method': 'sha256'
                    }

    def validate(self, username, password):
        # 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()



if __name__ == '__main__':

    import argparse
    parser = argparse.ArgumentParser()

    parser.add_argument('--file', default='passwords.json')

    subparsers = parser.add_subparsers(dest='cmd')

    add_parser = subparsers.add_parser('add')
    add_parser.add_argument('username')
    add_parser.add_argument('password')

    val_parser = subparsers.add_parser('validate')
    val_parser.add_argument('username')
    val_parser.add_argument('password')

    args = parser.parse_args()


    passwords = Passwords(args.file)
    if args.cmd == 'add':
        passwords.add(args.username, args.password)
        passwords.save()
    elif args.cmd == 'validate':
        print(passwords.validate(args.username, args.password))
    else:
        parser.print_help()