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, key):
    """
    Szyfrowanie i zapis pliku o nazwie filename (str) kluczem key (bytes)
    """
    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(filename, "wb") as file:
        file.write(encrypted_data)

def decrypt(filename, key):
    """
    Odszyfrowywanie i zapis pliku o nazwie filename (str). Plik zaszyfrowno kluczem key (bytes)
    """
    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(filename, "wb") as file:
        file.write(decrypted_data)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Prosty skrypt szyfrujący pliki.")
    parser.add_argument("file", help="Nazwa pliku do zaszyfrowania/odszyfrowania.")
    parser.add_argument("-g", "--generate-key", dest="generate_key", action="store_true",
                        help="Flaga określająca, czy klucz ma zostać wygenerowany (jeżeli nie, to użyty zostanie istniejący klucz).")
    parser.add_argument("-e", "--encrypt", action="store_true",
                        help="Szyfrowanie. Dozwolone opcje to tylko -e i -d.")
    parser.add_argument("-d", "--decrypt", action="store_true",
                        help="Szyfrowanie. Dozwolone opcje to tylko -e i -d.")

    args = parser.parse_args()
    file = args.file
    generate_key = args.generate_key

    if generate_key:
        write_key()
    # Wczytywanie klucza
    key = load_key()

    encrypt_ = args.encrypt
    decrypt_ = args.decrypt

    if encrypt_ and decrypt_:
        raise TypeError("Określ tryb pracy (szyfrowanie/odszyfrowywanie).")
    elif encrypt_:
        encrypt(file, key)
    elif decrypt_:
        decrypt(file, key)
    else:
        raise TypeError("Określić tryb pracy (szyfrowanie/odszyfrowywanie).")
