#!/usr/bin/env python

# salaus.py
# CT30A3700 Tietoturva
# Harjoitustyo
# Hanna Salopaasi
# 0358645
# hanna.salopaasi@lut.fi


import sys
import os
import re
import pickle
from Crypto.Cipher import AES
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA256
from Crypto import Random

def encrypt(input_file):
    file_to_crypt = open(input_file, 'r')
    output = open('salattusopimus.txt', 'w')
#Creating AES encryptor
    key = raw_input("Anna haluttu avain(1-32): ")
#Adding some padding to the key to make it long enough
    if len(key) <32:
        padding = '}'
        key = key + (32-len(key))*padding
    aes_encryptor = AES.new(key, AES.MODE_ECB)
#Creating signature
    rng = Random.new().read
    RSA_key = RSA.generate(1024, rng)
    text = file_to_crypt.read()  
    shahash = SHA256.new(text).digest()
    signature = RSA_key.sign(shahash, rng)
    rsa_publickey= RSA_key.publickey()
#Saving encrypted text and signature to a file called "salattusopimus.txt"
    output.write(str(signature))
    output.write('\n')
    output.write(aes_encryptor.encrypt(text))
#Saving the RSA public key to a file called "julkinenavain"
    key_file = open('julkinenavain', 'w')
    key_out = pickle.dump(rsa_publickey, key_file)
    file_to_crypt.close()
    key_file.close()

def decrypt(input_file, key_file):
    file_to_decrypt = open(input_file, 'r')
    key_file = open(key_file, 'r')
#Load RSA public key
    rsa_publickey = pickle.load(key_file)
#Read signature from the file and turn it into a tuple
    signature = eval(file_to_decrypt.readline())
#Add the same padding to the given decrypt key
    key = raw_input("Anna purkuavain: ")
    if len(key) <32:
        padding = '}'
        key = key + (32-len(key))*padding
    aes_decryptor = AES.new(key, AES.MODE_ECB)
    decrypted_file = aes_decryptor.decrypt(file_to_decrypt.read())
#Saving decrypted text to a file called "purettusopimus.txt"
    outfile = open('purettusopimus.txt', 'w')
    outfile.write(decrypted_file)
    print "Purettu tiedostoon: purettusopimus.txt"
#Calculating hash from decrypted text and verifying signature
    shahash = SHA256.new(decrypted_file).digest()
    if rsa_publickey.verify(shahash, signature):
        print "Allekirjoitus varmennettu."
    else:
        print "Virheellinen allekirjoitus tai purkuavain!"
    file_to_decrypt.close()
    outfile.close()


try:
    option = sys.argv[1]
    if option == "-salaa":
        encrypt(sys.argv[2])
    elif option == "-pura":
        decrypt(sys.argv[2], sys.argv[3])
    else:
        print "Virheellinen komento"
except IndexError:
    print "salaus.py -salaa tiedosto\nsalaus.py -pura tiedosto julkinenavain"