#!/bin/python

import os
import argparse
from hashlib import sha256
from base64 import b64encode, b64decode
from binascii import hexlify

def generate_hash(password, salt=None):
    if salt is None:
        salt = os.urandom(4)

    salted = salt + password.encode('utf-8')
    hashed = salt + sha256(salted).digest()
    return b64encode(hashed).decode('utf-8')

def test_hash(hashed, password, expected_salt=None):
    hash_bytes = b64decode(hashed)
    salt = hash_bytes[0:4]

    if expected_salt is not None and salt != expected_salt:
        return f'Hash {hexlify(hash_bytes)} with salt {hexlify(salt)} does not match expected salt {hexlify(expected_salt)}'

    if generate_hash(password, salt) == hashed:
        return f'Password matches hash'
    else:
        return f'Password does not match hash'


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Tool for hashing user passwords that can be added in RabbitMQ config files.')
    parser.add_argument('password', type=str, help='The password to hash or test against')
    parser.add_argument('-s', '--salt', dest='salt', type=str, help='Use a given salt instead of generatinga random one. If used with the --test argument, it verifies that the hash used this salt')
    parser.add_argument('-t', '--test', metavar='HASH', dest='test', type=str, help='Instead of hashing a password, check if a password matches a hash')
    args = parser.parse_args()

    if args.salt:
        salt = bytes.fromhex(args.salt)
    else:
        salt = None

    if args.test:
        print(test_hash(args.test, args.password, salt))
    else:
        print(generate_hash(args.password, salt))