keypair_generation.py

from nacl.signing import SigningKey, VerifyKey
import base64
import bcrypt

class HashValidationError(Exception):
    pass

class KeyPair(object):
    def __init__(self, private_key):
        # Save the keys.
        self.private_key = private_key

        # The private key is in fact a base64 encoded seed of a signing key.
        # The seed is used to generate a private signing key and a public key.
        # Given that from the seed all the other keys can be generated,
        # We treat the seed as a private key. The seed is 32 secure random bytes.
        self.seed = base64.b64decode(self.private_key)

        # Now we recreate the SigningKey object
        self.signing_key = SigningKey(seed=self.seed)
        # And the public key
        self.public_key = base64.b64encode(self.signing_key.verify_key.encode())

    @classmethod
    def generate(cls):
        # Generate a new random signing key
        signing_key = SigningKey.generate()

        # Save the seed, which we'll treat as the private key
        seed_bytes = signing_key._seed
        seed_base64 = base64.b64encode(seed_bytes)
        return cls(private_key=seed_base64)

    @classmethod
    def verify_msg(cls, message, public_key):
        # Create a VerifyKey object from the public key
        verify_key_bytes = base64.b64decode(public_key)
        verify_key = VerifyKey(verify_key_bytes)

        # Verify the decoded message
        message_bytes = base64.b64decode(message)
        return verify_key.verify(message_bytes)

    def sign_msg(self, message):
        # Sign a message with the signing key
        signed = self.signing_key.sign(bytes(message, encoding='utf8'))
        return base64.b64encode(signed)

    def hash(self):
        """
        Returns the bcrypt hash of the private key. This hash is what we should store in the database
        for a given user.
        :return: hash
        """
        return bcrypt.hashpw(self.private_key, bcrypt.gensalt())

    @classmethod
    def verify_hash(cls, private_key, hash):
        if not bcrypt.checkpw(private_key, hash):
            raise HashValidationError
        return cls(private_key)

    @classmethod
    def verify_hash_str(cls, private_key, hash):
        return cls.verify_hash(bytes(private_key, encoding="utf8"), bytes(hash, encoding="utf8"))

if __name__=="__main__":
    # This is what a regular flow would look like.
    # 1. Create a key pair
    user_id = input("Target user:")
    keypair = KeyPair.generate()

    # 2. In the database store the hash of the private key and the public key.
    user_keys = {}
    user_keys[user_id] = {
        "hash": str(keypair.hash(), "utf8"),
        "public_key": str(keypair.public_key, "utf8"), # In addition to the hash we store the public key.
    }
    print(f'\n\nThis is your private key (store in a secure location): {str(keypair.private_key, "utf8")}')
    print(user_keys)

    # 3. Now let's take the private key as an argument, in an API we'll grab this from
    #    the authorization header.
    input_private_key = input("Private key:")
    #   This would be the equivalent as fetching the hash from database
    hash = user_keys[user_id]["hash"]
    assert KeyPair.verify_hash_str(private_key=input_private_key, hash=hash)
    print(f"\nSuccesfully verified {user_id}")


    # The following are extra functionality for signing and verifying messages.
    # I don't expect we'll need these soon.
    # 4. Validate that the given keys produce signed messages
    msg = "test message"
    apikeys = KeyPair(input_private_key)
    signed = apikeys.sign_msg(msg)
    assert KeyPair.verify_msg(signed, apikeys.public_key), "Unable to verify message"
    print("The keys are valid and able to sign the messages (tried test msg: {})".format(msg))