from cryptography.fernet import Fernet
import os

def write_key():
    """
    Generowanie i zapis klucza do pliku
    """
    key = Fernet.generate_key()
    with open("klucz.key", "wb") as key_file:
        key_file.write(key)

def load_key():
    """
    Wczytanie pliku o nazwie `klucz.key` z bieżącego katalogu
    """
    return open("klucz.key", "rb").read()

def encrypt(filename, newfile, key):
    """
    Szyfrowanie obrazu filename (str) kluczem key (bytes). Zapis wyniku w pliku newfile (str).
    """
    f = Fernet(key)
    with open(filename, "rb") as file:
        # Wczytanie zawartości pliku
        file_data = file.read()
    # Szyfrowanie danych
    encrypted_data = f.encrypt(file_data)
    # Zapis zaszyfrowanego pliku
    with open(newfile, "wb") as file:
        file.write(encrypted_data)

def decrypt(filename, newfile, key):
    """
    Odszyfrowywanie obrazu filename (str), który zaszyfrowano kluczem key (bytes). Zapis wyniku w pliku newfile (str).
    """
    f = Fernet(key)
    with open(filename, "rb") as file:
        # Wczytanie zaszyfrowanej zawartości pliku
        encrypted_data = file.read()
    # Odszyfrowywanie danych
    decrypted_data = f.decrypt(encrypted_data)
    # Odtworzenie oryginalnego pliku
    with open(newfile, "wb") as file:
        file.write(decrypted_data)

key = Fernet.generate_key()

enc = encrypt("./r06_sekretny_obraz.jpg", "./e_r06_sekretny_obraz.jpg", key)
dec = decrypt("./e_r06_sekretny_obraz.jpg", "./d_r06_sekretny_obraz.jpg", key)