from pbkdf2 import PBKDF2
from hashlib import sha512
from base64 import b64encode
import argparse
import secrets
import string
import psycopg2


parser = argparse.ArgumentParser(description='genpw')
parser.add_argument('--length', '-l',
                    help='Length of auto-generated password', 
                    default='24',
                    required=False)
parser.add_argument('--password', '-p',
                    help='Password', 
                    required=False)
parser.add_argument('--username', '-u',
                    help='Username',
                    required=False)
parser.add_argument('--topic', '-t',
                    help='Initially granted topic',
                    required=False)
parser.add_argument('--acl', '-a',
                    help='ACL value for topic, Bit0=read, Bit1=write, Bit2=subscribe',
                    required=False)
parser.add_argument('--printonly', '-o',
                    help='Just print the password hash, do not write to database',
                    action='store_true')

args = parser.parse_args()
length = args.length
password = args.password

print_only = args.printonly

alphabet = string.ascii_letters + string.digits
iterations = 100000

if (not password):
    if (not length):
        raise Exception("Either length or password must be given")
    password = ''.join(secrets.choice(alphabet) for i in range(int(length)))

salt = secrets.token_bytes(16)
hash = b64encode(PBKDF2(password, salt, iterations=iterations, digestmodule=sha512).read(64)).decode()

salt_b64 = b64encode(salt).decode()

pw = f"PBKDF2$sha512${iterations}${salt_b64}${hash}"
print(f"{password=}")
print(f"hash={pw}")

if not print_only:
    login = args.username
    if (not login):
        raise Exception("For writing to database a username must be given")
    topic = args.topic
    acl = args.acl

    conn = psycopg2.connect()
    conn.autocommit = False

    try:
        with conn:
            with conn.cursor() as cur:
                cur.execute("""
                            insert into users_t (username, pw) 
                                values(%(username)s, %(pw)s) 
                                on conflict on constraint users_t_uk_username 
                                    do update set pw = %(pw)s
                                returning id 
                            """,
                            { 'username': login, 'pw': pw })
                res = cur.fetchone()
                if res is None:
                    raise Exception("Unable to add user to database")
                id = res[0]
                print("User added to database")
            if (topic and acl):
                acl = int(acl)
                with conn.cursor() as cur:
                    cur.execute('insert into acls_t ("user", topic, rw) values(%(user)s, %(topic)s, %(rw)s)',
                                { 'user': id, 'topic': topic, 'rw': acl })
                    print("ACL added to database")
    finally:
        if conn:
            conn.close()