import rsa
import sys
import datetime

def generate_keys():
    (pubKey, privKey) = rsa.newkeys(1024)
    with open('rsa-keys/publickey.pem', 'wb') as f:
        f.write(pubKey.save_pkcs1('PEM'))

    with open('rsa-keys/privatekey.pem', 'wb') as f:
        f.write(privKey.save_pkcs1('PEM'))

def load_keys():
    with open('rsa-keys/publickey.pem', 'rb') as f:
        pubKey = rsa.PublicKey.load_pkcs1(f.read())

    with open('rsa-keys/privatekey.pem', 'rb') as f:
        privKey = rsa.PrivateKey.load_pkcs1(f.read())

    return pubKey, privKey

def encrypt(msg, key):
    return rsa.encrypt(msg.encode('ascii'), key)

def decrypt(ciphertext, key):
    try:
        return rsa.decrypt(ciphertext, key).decode('ascii')
    except:
        return False

def sign_sha1(msg, key):
    return rsa.sign(msg.encode('ascii'), key, 'SHA-1')

def verify_sha1(msg, signature, key):
    try:
        return rsa.verify(msg.encode('ascii'), signature, key) == 'SHA-1'
    except:
        return False

generate_keys()
pubKey, privKey = load_keys()

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

n = 120
chunks = [text[i:i+n] for i in range(0, len(text), n)]
ciphertext = bytes(''.encode('ascii'))
signature = bytes(''.encode('ascii'))
encryptdelta = int(filereaddelta.total_seconds() * 1000)
decryptdelta = 0
filename = filename.replace(".", "_encrypted.")
filewrite = open(filename, 'w')
fileread = open(filename, 'r')
file.close()
for chunk in chunks:
    start_time = datetime.datetime.now()
    ciphertext = encrypt(chunk, pubKey)
    signature = sign_sha1(chunk, privKey)
    filewrite.write(str(ciphertext))
    delta = datetime.datetime.now() - start_time
    encryptdelta += int(delta.total_seconds() * 1000)
    start_time = datetime.datetime.now()
    text = fileread.readline()
    plaintext = decrypt(ciphertext, privKey)
    delta = datetime.datetime.now() - start_time
    decryptdelta += int(delta.total_seconds() * 1000)

filewrite.close()
fileread.close()
print('Šifrování trvalo: %s ms' % encryptdelta)
print('Rozšifrování trvalo: %s ms' % decryptdelta)