Python RSA tools

rsa_sign.py

# Run `pip install pycryptodome` to install Crypto library
from Crypto.PublicKey import RSA
from Crypto.Signature import pss
from Crypto.Hash import SHA256
import base64

def sign():
    with open("private.key", "r") as file:
        data = file.read()
    private_key = RSA.import_key(data)
    msg = input("Message: ").encode()
    hashed = SHA256.new(msg)
    signature = pss.new(private_key).sign(hashed)
    with open("signed_message.txt", "w") as output:
        output.write(msg.decode())
        output.write("\n--\n")
        output.write(base64.b64encode(signature).decode())
    # print(base64.b64encode(signature).decode())

def verify():
    keyfile = input("Public key: ")
    with open(keyfile, "r") as f:
        data = f.read()
    public_key = RSA.import_key(data)
    msg = input("Message: ").encode()
    signature = input("Signature: ").encode()
    signature = base64.b64decode(signature)
    hashed = SHA256.new(msg)
    try:
        pss.new(public_key).verify(hashed, signature)
        print("Verified")
    except:
        print("Verification failed")

choice = input("(S)ign or (V)erify: ").lower()
if choice == "s":
    sign()
elif choice == "v":
    verify()