python compute a hash of file and sign it with cryptography.hazmat

example-hash.py

import os
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.serialization import (
    Encoding, PrivateFormat, PublicFormat, NoEncryption
)

def generate_key_pair():
    """Generate an RSA key pair."""
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
    )
    public_key = private_key.public_key()
    
    # Serialize keys to PEM format
    private_pem = private_key.private_bytes(
        encoding=Encoding.PEM,
        format=PrivateFormat.PKCS8,
        encryption_algorithm=NoEncryption()
    )
    
    public_pem = public_key.public_bytes(
        encoding=Encoding.PEM,
        format=PublicFormat.SubjectPublicKeyInfo
    )
    
    # Save keys to files
    with open("private_key_random.pem", "wb") as f:
        f.write(private_pem)
    
    with open("public_key_random.pem", "wb") as f:
        f.write(public_pem)
    
    return private_key, public_key

def compute_file_hash(file_path):
    """Compute SHA-256 hash of a file."""
    digest = hashes.Hash(hashes.SHA256())
    
    with open(file_path, "rb") as f:
        # Read file in chunks to handle large files
        chunk_size = 8192
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            digest.update(chunk)
    
    return digest.finalize()

def sign_hash(private_key, file_hash):
    """Sign a hash using RSA-PSS."""
    signature = private_key.sign(
        file_hash,
        padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH
        ),
        hashes.SHA256()
    )
    return signature

def verify_signature(public_key, file_hash, signature):
    """Verify the signature of a hash."""
    try:
        public_key.verify(
            signature,
            file_hash,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        return True
    except Exception:
        return False

def main():
    # File to hash and sign
    file_path = input("Enter the path to the file: ")
    
    if not os.path.exists(file_path):
        print(f"Error: File '{file_path}' not found.")
        return
    
    # Generate key pair
    print("Generating RSA key pair...")
    private_key, public_key = generate_key_pair()
    print("Keys generated and saved to 'private_key.pem' and 'public_key.pem'")
    
    # Compute file hash
    print("Computing file hash...")
    file_hash = compute_file_hash(file_path)
    print(f"SHA-256 hash: {file_hash.hex()}")
    
    # Sign the hash
    print("Signing the hash...")
    signature = sign_hash(private_key, file_hash)
    
    # Save signature to file
    with open("signature.bin", "wb") as f:
        f.write(signature)
    print(f"Signature saved to 'signature.bin'")
    
    # Verify the signature
    print("Verifying signature...")
    if verify_signature(public_key, file_hash, signature):
        print("Signature verified successfully!")
    else:
        print("Signature verification failed!")

if __name__ == "__main__":
    main()