from Crypto.Cipher import AES
from Crypto.Hash import SHA256
import os, random, struct
from Crypto import Random

def encrypt_file(key, filename):
    chunk_size = 64 * 1024
    output_filename = filename + '.encrypted'
    # Losowo wygenerowany wektor inicjujący.
    iv = Random.new().read(AES.block_size)
    # Zdefiniowanie szyfru.
    encryptor = AES.new(key, AES.MODE_CBC, iv)
    # Określenie wielkości pliku.
    filesize = os.path.getsize(filename)
    # Otwarcie pliku wyjściowego i zapisanie w nim wielkości pliku.
    # Wykorzystywany jest w tym celu moduł struct.
    with open(filename, 'rb') as inputfile:
        with open(output_filename, 'wb') as outputfile:
            outputfile.write(struct.pack('<Q', filesize))
            outputfile.write(iv)
            while True:
                chunk = inputfile.read(chunk_size)
                if len(chunk) == 0:
                    break
                elif len(chunk) % 16 != 0:
                    chunk += bytes(' ', 'utf-8') * (16 - len(chunk) % 16)
                outputfile.write(encryptor.encrypt(chunk))

def decrypt_file(key, filename):
    chunk_size = 64 * 1024
    output_filename = os.path.splitext(filename)[0]
    # Otwarcie zaszyfrowanego pliku i odczytanie z niego wielkości oryginalnego pliku i wektora inicjującego.
    # Wektor jest niezbędny do zdefiniowania szyfru.
    with open(filename, 'rb') as infile:
        origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
        iv = infile.read(16)
        # Zdefiniowanie szyfru przy użyciu klucza i wektora.
        decryptor = AES.new(key, AES.MODE_CBC, iv)
        # Odszyfrowane dane są zapisywane w pliku weryfikacyjnym.
        # Porównując uzyskany plik z oryginalnym, można sprawdzić
        # poprawność szyfrowania.
        with open(output_filename, 'wb') as outfile:
            while True:
                chunk = infile.read(chunk_size)
                if len(chunk) == 0:
                    break
                outfile.write(decryptor.decrypt(chunk))
            outfile.truncate(origsize)

def getKey(password):
    hasher = SHA256.new(password)
    return hasher.digest()

def main():
    choice = input("Wybierz operację: (S)zyfrowanie, (D)eszyfrowanie: ")
    if choice == 'S':
        filename = input('Podaj nazwę pliku do zaszyfrowania: ')
        password = input('Podaj hasło: ')
        encrypt_file(getKey(password.encode("utf8")), filename)
        print('Koniec.')
    elif choice == 'D':
        filename = input('Podaj nazwę pliku do odszyfrowania: ')
        password = input('Podaj hasło: ')
        decrypt_file(getKey(password.encode("utf8")), filename)
        print('Koniec.')
    else:
        print('Błędny wybór.')

if __name__ == "__main__":
    main()