from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.exceptions import InvalidSignature

# Generate RSA key pair
priv = rsa.generate_private_key(public_exponent=65537, key_size=2048)
pub = priv.public_key()

# Message
msg = b"Important message to be signed."

# Sign message
sig = priv.sign(msg,
    padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
    hashes.SHA256())
print("Signature created.")

# Verify signature
try:
    pub.verify(sig, msg,
        padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
        hashes.SHA256())
    print("Signature verified.")
except InvalidSignature:
    print("Verification failed!")

# Export public key (PEM)
pem = pub.public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)
print("\nPublic Key:\n", pem.decode())
