from Crypto.Cipher import AES
from secrets import token_bytes
import sys
import datetime

def encrypt(msg):
    cipher = AES.new(key, AES.MODE_EAX)
    nonce = cipher.nonce
    ciphertext, tag = cipher.encrypt_and_digest(msg.encode('ascii'))
    return nonce, ciphertext, tag

def decrypt(nonce, ciphertext, tag):
    cipher = AES.new(key, AES.MODE_EAX, nonce=nonce)
    plaintext = cipher.decrypt(ciphertext)
    try:
        cipher.verify(tag)
        return plaintext.decode('ascii')
    except:
        return False

key = token_bytes(16)

print('Zahajuji čtení souboru.')
start_time = datetime.datetime.now()
filename = str(sys.argv[1])
file = open(filename, 'r')
lines = file.readlines()
file.close()
text = ''
for line in lines:
    text += str(line)
filereaddelta = datetime.datetime.now() - start_time
print('Čtení souboru úspěšné.')


start_time = datetime.datetime.now()
nonce, ciphertext, tag = encrypt(text)
filename = filename.replace(".", "_encrypted.")
file = open(filename, 'w')
file.write(str(ciphertext))
file.close()
delta = datetime.datetime.now() - start_time + filereaddelta
print('Šifrování trvalo: %s ms' % int(delta.total_seconds() * 1000))

start_time = datetime.datetime.now()
file = open(filename, 'r')
lines = file.readlines()
file.close()
text = ''
for line in lines:
    text += str(line)
plaintext = decrypt(nonce, ciphertext, tag)
delta = datetime.datetime.now() - start_time
print('Rozšifrování trvalo: %s ms' % int(delta.total_seconds() * 1000))