import zlib
import base64
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
from pathlib import Path

# Funkcja generująca nową parę kluczy
def generate_new_key_pair():
    # Generowanie pary klucz publiczny – klucz prywatny; długość klucza to 4096 bitów (512 bajtów)    
    new_key = RSA.generate(4096, e=65537)

    # Klucz prywatny w formacie PEM
    private_key = new_key.exportKey("PEM")

    # Klucz publiczny w formacie PEM
    public_key = new_key.publickey().exportKey("PEM")

    private_key_path = Path('private.pem')
    private_key_path.touch(mode=0o600)
    private_key_path.write_bytes(private_key)

    public_key_path = Path('public.pem')
    public_key_path.touch(mode=0o664)
    public_key_path.write_bytes(public_key)


# Funkcja szyfrująca RSA
def encrypt_blob(blob, public_key):
    # Importowanie klucza publicznego, szyfrowanie za pomocą PKCS1_OAEP
    rsa_key = RSA.importKey(public_key)
    rsa_key = PKCS1_OAEP.new(rsa_key)

    # Na początek kompresja danych
    blob = zlib.compress(blob)
    # W przypadku PKCS1_OAEP długość pojedynczego fragmentu/bloku (chunk_size)
    # to długość klucza w bajtach minus 42. Dane są szyfrowane w blokach/fragmentach
    chunk_size = 470
    offset = 0
    end_loop = False
    encrypted = bytearray()

    while not end_loop:
        # Fragment/blok
        chunk = blob[offset:offset + chunk_size]

        # Jeżeli dany blok/fragment jest krótszy niż chunk_size, to dopełnia się go znakami " "
        # Oznacza to, że osiągnięty został rozmiar pliku i pętla może zostać przerwana
        if len(chunk) % chunk_size != 0:
            end_loop = True
            # chunk += b" " * (chunk_size - len(chunk))
            chunk += bytes(chunk_size - len(chunk))
        # Dodawanie zaszyfrowanego fragmentu do reszty
        encrypted += rsa_key.encrypt(chunk)

        # Zwiększanie offsetu o rozmiar fragmentu/bloku
        offset += chunk_size

    # Kodowanie rezultatu w base64
    return base64.b64encode(encrypted)

# Funkcja odszyfrowująca RSA
def decrypt_blob(encrypted_blob, private_key):

    # Importowanie klucza prywatnego, który jest niezbędny do odszyfrowania za pomocą PKCS1_OAEP
    rsakey = RSA.importKey(private_key)
    rsakey = PKCS1_OAEP.new(rsakey)

    # Dekodowanie base64
    encrypted_blob = base64.b64decode(encrypted_blob)

    # Wartość chunk_size będzie równa długości klucza prywatnego w bajtach
    # Odszyfrowywanie będzie przebiegać we fragmentach/blokach
    chunk_size = 512
    offset = 0
    decrypted = bytearray()

    # Pętla trwa, dopóki są jeszcze bloki/fragmenty do odszyfrowania
    while offset < len(encrypted_blob):
        # Fragment/blok
        chunk = encrypted_blob[offset: offset + chunk_size]

        # Dodawanie odszyfrowanego fragmentu/bloku do reszty
        decrypted += rsakey.decrypt(chunk)

        # Zwiększanie offsetu o rozmiar fragmentu/bloku
        offset += chunk_size

    # Funkcja zwraca odkodowane dane po dekompresji
    return zlib.decompress(decrypted)


# generate_new_key_pair() # Odkomentuj, jeżeli nie masz jeszcze pary kluczy

print("Program szuka pliku 'chmurka.jpg' w katalogu 'Rozdział 8'.")

private_key = open('klucz_prywatny.pem').read()

print("Wczytano klucz prywatny.")

public_key = open('klucz_publiczny.pem').read()

print("Wczytano klucz publiczny.")

unencrypted_file = Path('./chmurka.jpg')
encrypted_file = unencrypted_file.with_suffix('.dat')
encrypted_blob = encrypt_blob(unencrypted_file.read_bytes(), public_key)

print("Plik został zaszyfrowany.")

# Zapis zaszyfrowanej postaci do pliku
fd = open("./e_chmurka.jpg", "wb")
fd.write(encrypted_blob)
fd.close()

print("Zaszyfrowany plik to e_chmurka.jpg")

# Odszyfrowywanie pliku
fd = open("./e_chmurka.jpg", "r")
encrypted_blob = fd.read()
fd.close()

print()
print("Zawartość zaszyfrowanego pliku jest zbyt długa, aby ją wyświetlić.")
print()

# Zapis odszyfrowanych danych do pliku
fd = open("./d_chmurka.jpg", "wb")
fd.write(decrypt_blob(encrypted_blob, private_key))
fd.close()

# Odszyfrowywanie danych
decrypt_blob(encrypted_blob, private_key)

print("Plik został odszyfrowany. Rezultat umieszczono w pliku  d_chmurka.jpg.")

