Python, using the cryptopgraphy package: Create public key and private key, save the keys, load the keys, sign a message, and verify a message

cryptography_signing_helper.py

from cryptography.hazmat.primitives import serialization as crypto_serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.backends import default_backend as crypto_default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import serialization
from cryptography.exceptions import InvalidSignature
import time
import json


def read_keypair(from_file_path: str):
    with open(from_file_path, "r") as fp:
        info = json.loads(fp.read())
        return info["public_key"], info["private_key"]


def save_keypair(to_file_path: str, public_key: bytes, private_key: bytes):
    with open(to_file_path, "w") as fp:
        json.dump(
            {
                "public_key": public_key.decode(),
                "private_key": private_key.decode()
            },
            fp
        )


def generate_keypair():

    key = rsa.generate_private_key(
        backend=crypto_default_backend(),
        public_exponent=65537,
        key_size=2048
    )

    private_key = key.private_bytes(
        crypto_serialization.Encoding.PEM,
        crypto_serialization.PrivateFormat.PKCS8,
        crypto_serialization.NoEncryption()
    )

    public_key = key.public_key().public_bytes(
        crypto_serialization.Encoding.OpenSSH,
        crypto_serialization.PublicFormat.OpenSSH
    )

    return public_key, private_key


def sign_message(message: bytes, private_key: bytes):
    private_key = serialization.load_pem_private_key(
        private_key,
        password=None
    )
    signature = private_key.sign(
        message,
        padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH
        ),
        hashes.SHA256()
    )
    return signature


def verify_message(message: bytes, signature: bytes, public_key: bytes):
    public_key = serialization.load_ssh_public_key(public_key)
    try:
        public_key.verify(
            signature,
            message,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        return True
    except InvalidSignature:
        return False


def main():

    # Generate keypair
    file_path = "keypair.json"
    public_key, private_key = generate_keypair()
    save_keypair(file_path, public_key, private_key)
    print(public_key)
    print(private_key)

    # Read the keypair to ensure it saved properly
    public_key_read, private_key_read = read_keypair(file_path)
    print(public_key_read)
    print(private_key_read)

    # Generate message and sign it
    message = str(int(time.time())).encode()
    signature = sign_message(message, private_key_read.encode())
    print(signature)

    # Verify it
    is_verified = verify_message(message=message, signature=signature, public_key=public_key_read.encode())
    print("Does it match? (Should be True): ", is_verified)

    # Attempt to verify it with a different public key
    second_public_key, _ = generate_keypair()
    is_verified = verify_message(message=message, signature=signature, public_key=second_public_key)
    print("Does it match? (Should be False): ", is_verified)

    # Send message and signature as JSON payload
    print(signature)
    payload = {
        "message": message.decode(),
        "signature": list(signature)
    }
    print(payload)

    signature = bytes(payload['signature'])
    print(signature)

    print(f"Payload to send: {json.dumps(payload)}")

    # Now send this payload using the requests package.
    # For you to do...


if __name__ == '__main__':
    main()