from Crypto.PublicKey import RSA 
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Hash import SHA256
from Crypto.Signature import PKCS1_v1_5

def generate(bit_size):
    keys = RSA.generate(bit_size)
    return keys

def encrypt(pub_key, data):
    cipher = PKCS1_OAEP.new(pub_key)
    return cipher.encrypt(data)

def decrypt(priv_key, data):
    cipher = PKCS1_OAEP.new(priv_key)
    return cipher.decrypt(data)

keys = generate(2048)
print("Klucz publiczny:")
print(keys.publickey().export_key('PEM').decode(), end='\n\n')
with open("public.key", 'wb') as file:
    file.write(keys.publickey().export_key())
print("Klucz prywatny:")
print(keys.export_key('PEM').decode())
with open("private.key", 'wb') as file:
    file.write(keys.export_key('PEM'))
text2cipher = "Poufny komunikat".encode("utf8")
hasher = SHA256.new(text2cipher)
signer = PKCS1_v1_5.new(keys)
signature = signer.sign(hasher)
verifier = PKCS1_v1_5.new(keys)
if verifier.verify(hasher, signature):
    print('Podpis poprawny!')
else:
    print('Komunikat został podpisany innym kluczem lub został zmodyfikowany!')
encrypted_data = encrypt(keys.publickey(),text2cipher)
print("Zaszyfrowany tekst:", encrypted_data)
decrypted_data = decrypt(keys,encrypted_data)
print("Odszyfrowany tekst:", decrypted_data.decode())