#!/usr/bin/python

from __future__ import print_function
import sys
from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA
from Crypto.Cipher import AES
import binascii


class Signature:
    pvt_key = ''
    filename = ''
    hash = SHA256.new()

    def __init__(self, pvt, file):
        self.pvt_key = pvt
        self.filename = file

    '''
    @params: N/A
    @precondition: A valid file_path must have been provided when instantiating the class.
    @postcondition: A SHA256 hash of the file has been returned.
    @short-description: Using a SHA256 object, given the file_path, the function opens the file at the specified path,
    updates the hash object and returns the hash of the file after closing the file.
    '''
    def hashFile(self):
        fd = open(self.filename, 'rb')
        self.hash.update(fd.read())
        fd.close()
        return self.hash.digest()

    '''
    @params: N/A
    @precondition: A valid path to a file containing the private key has been provided upon creating the object.
    @postcondition: A signature in hex format of the file has been returned.
    @short-description: An RSA_Object is created and it's key is set to the provided private key, using the _RSAObj
    the hashed data from the file is signed. The sign() method in the RSA module of the PyCrypto library returns a tuple
    of two items, the first containing the signature in long format. We therefore convert it to hex and remove the trailing
    'L' from the long and return the signature after having closed the file.
    '''
    def signFile(self):
        fd = open(self.pvt_key, 'rb')
        key = RSA.importKey(fd.read())
        signature = key.sign(self.hash.digest(), '')
        fd.close()
        return hex(signature[0]).rstrip('L')

    '''
    @params: the path to the file containing the public key and the path to the file containing the signature.
    @precondition: A valid path to both a public key and a signature have been provided upon object creation.
    @postcondition: A bool is returned:
                    -TRUE- if the verification was successful.
                    -FALSE- if the verification was not successful.
    @short-description: The file is hashed and using an _RSAObj, after setting the public key and the signature, the 
    verify() method of the RSA module in the PyCrypto accepts as arguments a hashed file and a signature in the format
    of the return value of the sing() function, therefore the signature has to be reconverted to a long format and the
    tuple has to be reconstructed with None as the second item.
    '''
    def verifyFile(self, public, signature):
        s = open(signature, 'rb')
        p = open(public, 'rb')
        hashed = self.hashFile()
        pub_key = RSA.importKey(p.read())
        verify = pub_key.verify(hashed, (long(s.read(), 16), None))
        p.close()
        s.close()
        return verify

class Encryption:
    init_vector = ''
    filename = ''
    secret_key = ''
    padded = ''
    AES.block_size = 16

    def __init__(self, iv, file, key):
        self.filename = file
        self.init_vector = iv
        self.secret_key = SHA256.new(key).digest()

    '''
    @params: N/A
    @precondition: A valid file path containing the plain text has been specified upon class instantiation.
    @postcondition: A string containing the text padded using PKCS7 padding has been returned in order to get a length
    of the plaintext to be a multiple of 16 for AES.
    @short-description: The file in plain text is opened and read into a string (plain), according to PKCS7 standards
    it is padded if the length is not a multiple of 16.
    '''
    def padPlain(self):
        fd = open(self.filename, 'rb')
        plain = fd.read()
        fd.close()
        pkcs7_padding = (AES.block_size - (len(plain) % AES.block_size))
        if len(plain) % 16 != 0:
            self.padded = plain + chr(pkcs7_padding) * pkcs7_padding
        return self.padded

    '''
    @params: N/A
    @precondition: A valid file path containing the plain text has been specified upon class instantiation.
    @postcondition: The encrypted plaintext has been returned
    @short-description: Using an AES object initialised with the secret key and an initialisation vector, set to CBC mode,
    the encrypted plain text is generated by calling the encrypt() member function in the AES module of PyCrypto.
    '''
    def encrypt(self):
        cipher = AES.new(self.secret_key, AES.MODE_CBC, self.init_vector)
        encrypted = self.init_vector + cipher.encrypt(self.padPlain())
        return encrypted

    '''
    @params: A path file to the encrypted text.
    @precondition: A valid path to the encrypted file has been specified.
    @postcondition: The decrypted text has been returned.
    @short-description: The encrypted file is opened and read into a string (enc), an AES object is created with the 
    secret key and the initialisation vector in CBC mode. The decrypt() member function of the AES module of PyCrypto
    is used to decrypt the enc string (with the iv trimmed off). The decrypted text is then unpadded by finding the
    non-unicode chars at the end of the decrypted string (PKCS7 standard) in the second
    to last line of the suntionand the plain string is finally returned.
    '''
    def decrypt(self, file):
        fd = open(file, 'rb')
        enc = fd.read()
        fd.close()
        cipher = AES.new(self.secret_key, AES.MODE_CBC, self.init_vector)
        decrypted = cipher.decrypt(enc[AES.block_size:])
        plain = decrypted[:-ord(decrypted[-1:])]
        return plain


if len(sys.argv) not in (4,5):
    print("Error: Possible Usages:\n"
          "./vault.py -s <path_to_file> <path_to_private_key>\n"
          "./vault.py -v <path_to_file> <path_to_public_key> <path_to_signature>\n"
          "./vault.py -e <path_to_file> <secret_key> <iv>\n"
          "./vault.py -d <path_to_file> <secret_key> <iv>")

file_path = sys.argv[2]

if sys.argv[1] == '-s':
    pvt_key = sys.argv[3]
    sign = Signature(pvt_key, file_path)
    hashed = sign.hashFile()
    signature = sign.signFile()
    print(signature)
elif sys.argv[1] == '-v':
    pub_key = sys.argv[3]
    signature = sys.argv[4]
    sign = Signature(None, file_path)
    verification = sign.verifyFile(pub_key, signature)
    if verification:
        sys.exit(0)
    else:
        sys.exit(1)
elif sys.argv[1] == '-e':
    if sys.argv[3][:2] == '0x':
        secret_key = sys.argv[3][2:]
    else:
        secret_key = sys.argv[3]

    if sys.argv[4][:2] == '0x':
        init_vector = binascii.unhexlify(sys.argv[4][2:])
    else:
        init_vector = binascii.unhexlify(sys.argv[4])

    enc = Encryption(init_vector, file_path, secret_key)
    encrypted = enc.encrypt()
    print(encrypted, end="")
elif sys.argv[1] == '-d':
    if sys.argv[3][:2] == '0x':
        secret_key = sys.argv[3][2:]
    else:
        secret_key = sys.argv[3]

    if sys.argv[4][:2] == '0x':
        init_vector = binascii.unhexlify(sys.argv[4][2:])
    else:
        init_vector = binascii.unhexlify(sys.argv[4])

    dec = Encryption(init_vector, file_path, secret_key)
    decrypted = dec.decrypt(file_path)
    print(decrypted, end="")
else:
    print("Error: Unknown Command!")
