Python RSA encryption and decryption

my-python-rsa.py

import base64
import rsa # python3 -m pip install rsa
import os

def save_key(key: rsa.PublicKey | rsa.PrivateKey) -> bytes:
    return key.save_pkcs1("PEM")

def load_key(data: bytes, key_type: rsa.PublicKey | rsa.PrivateKey) -> rsa.PublicKey | rsa.PrivateKey:
    return key_type.load_pkcs1(data)

def generate_rsa_keys(length: int):
    if length < 90:
        raise ValueError(f"\"length\" must be superior to 90 but it's {length}")
    return rsa.newkeys(length)

def encrypt(data: bytes, public_key: rsa.PublicKey) -> bytes:
    chunk_size = rsa.common.byte_size(public_key.n) - 11
    splitted_message = [data[i:i+chunk_size] for i in range(0, len(data), chunk_size)]
    encrypted_message = [base64.b64encode(rsa.encrypt(chunk, public_key)) for chunk in splitted_message]
    return b';'.join(encrypted_message)

def decrypt(data: bytes, private_key: rsa.PrivateKey) -> bytes:
    if data == b"":
        return b""
    chunk_size = rsa.common.byte_size(private_key.n) - 11
    splitted_message = data.split(b';')
    if any(len(chunk) != chunk_size for chunk in splitted_message[:-1]):
        raise ValueError("Invalid chunk size !")
    return b''.join([rsa.decrypt(base64.b64decode(chunk), private_key) for chunk in splitted_message])

if __name__ == "__main__":
    print("+---------------------------------------+")
    print("+ Generating 2048 bits rsa key          +")
    pub, priv = generate_rsa_keys(2048)
    print("+ Key generated                         +")
    print("+ Saving public key to \"id_rsa.pub\"     +")
    with open("id_rsa.pub", 'wb+') as f:
        f.write(save_key(pub))
    print("+ Public key saved                      +")
    print("+ Saving private key to \"id_rsa.priv\"   +")
    with open("id_rsa.priv", 'wb+') as f:
        f.write(save_key(priv))
    print("+ Private key saved                     +")
    print("+ Deleting old keys                     +")
    del pub
    del priv
    print("+ Old keys deleted                      +")
    print("+ Loading public key                    +")
    with open("id_rsa.pub", 'rb') as f:
        pub = load_key(f.read(), rsa.PublicKey)
    print("+ Public key loaded                     +")
    print("+ Loading private key                   +")
    with open("id_rsa.priv", 'rb') as f:
        priv = load_key(f.read(), rsa.PrivateKey)
    print("+ Private key loaded                    +")
    print("+ Generating encrypted message          +")
    
    text = "Hello world, encrypted !"
    import sys
    if len(sys.argv) >= 2:
        text = ' '.join(sys.argv[1:])
    text = text.encode()
    encrypted = encrypt(text, pub)
    del text
    print("+ Encrypted text generated              +")
    print(f"+ {encrypted.decode()}")
    print("+ Decoding message                      +")
    decoded = decrypt(encrypted, priv)
    del encrypted
    print("+ Message decoded                       +")
    print(f"+ {decoded}")
    
    # Remove RSA keys (comment to keep them)
    if os.path.isfile("id_rsa.pub"):
        os.remove("id_rsa.pub")
    if os.path.isfile("id_rsa.priv"):
        os.remove("id_rsa.priv")