from Crypto.Cipher import DES3, DES
from Crypto import Random
import time


class DES3Crypto(object):
        def __init__(self):
            self.block_size = 8

        def __getPadding(self, data):
            pad_len = self.block_size - (len(data) % self.block_size)
            return data + bytes([pad_len] * pad_len)

        def __rmPadding(self, data):
            pad_len = data[-1]
            return data[:-pad_len]

        def __innernCBCtripleDES(self, data, key, iv, crypt='ENCRYPT'):
            result = b''
            iv1 = iv2 = iv3 = iv
            key1 = key[:8]
            key2 = key[8:16]
            if len(key) == 16:
                key3 = key1
            else:
                key3 = key[16:]
            for x in range(0, len(data), self.block_size):
                if crypt == 'ENCRYPT':
                    cipher = DES.new(key1, DES.MODE_CBC, iv1)
                    block = cipher.encrypt(data[x:x+self.block_size])
                    iv1 = block
                    cipher = DES.new(key2, DES.MODE_CBC, iv2)
                    block = cipher.decrypt(block)
                    iv2 = block
                    cipher = DES.new(key3, DES.MODE_CBC, iv3)
                    block = cipher.encrypt(block)
                    iv3 = block
                    result += block
                else:
                    cipher = DES.new(key1, DES.MODE_CBC, iv1)
                    block = cipher.decrypt(data[x:x+self.block_size])
                    iv1 = data[x:x+8]
                    cipher = DES.new(key2, DES.MODE_CBC, iv2)
                    iv2 = block
                    block = cipher.encrypt(block)
                    cipher = DES.new(key3, DES.MODE_CBC, iv3)
                    iv3 = block
                    block = cipher.decrypt(block)
                    result += block
            return result

        def __padTripleDES(self, data, key, crypt='ENCRYPT'):
            key1 = key[:8]
            key2 = key[8:16]
            if len(key) == 16:
                key3 = key1
            else:
                key3 = key[16:]
            if crypt == 'ENCRYPT':
                cipher = DES.new(key1, DES.MODE_ECB)
                data = cipher.encrypt(data)
                padding = Random.new().read(DES3.block_size // 2)
                data = self.__getPadding(padding + data)
                cipher = DES.new(key2, DES.MODE_ECB)
                data = cipher.decrypt(data)
                padding = Random.new().read(DES3.block_size // 2)
                data = self.__getPadding(padding + data)
                cipher = DES.new(key3, DES.MODE_ECB)
                data = cipher.decrypt(data)
            else:
                cipher = DES.new(key3, DES.MODE_ECB)
                data = cipher.decrypt(data)
                data = self.__rmPadding(data[self.block_size//2:])
                cipher = DES.new(key2, DES.MODE_ECB)
                data = cipher.encrypt(data)
                data = self.__rmPadding(data[self.block_size//2:])
                cipher = DES.new(key1, DES.MODE_ECB)
                data = cipher.decrypt(data)
            return data

        def getMode(self):
            print('1 - ECB, 2 - CBC, 3 - CFB,\n4 - PGP, 5 - OFB,'
                  ' 6 - CTR,\n7 - OPENPGP, 8 - innerCBC, 9 - ECBpad')

        def tripleDESencrypt(self, data, key, mode, iv):
            data = self.__getPadding(data)
            if mode == 8:
                encrypt = self.__innernCBCtripleDES(data, key, iv, 'ENCRYPT')
            elif mode == 9:
                encrypt = self.__padTripleDES(data, key, 'ENCRYPT')
            else:
                cipher = DES3.new(key, mode, iv)
                encrypt = cipher.encrypt(data)
            return encrypt

        def tripleDESdecrypt(self, encrypt, key, mode, iv):
            if mode == 8:
                data = self.__innernCBCtripleDES(encrypt, key, iv, 'DECRYPT')
            elif mode == 9:
                data = self.__padTripleDES(encrypt, key, 'DECRYPT')
            else:
                cipher = DES3.new(key, mode, iv)
                data = cipher.decrypt(encrypt)
            return self.__rmPadding(data)


        def testTime(self, data, key, iv):
            start = time.time()
            encrypt = self.tripleDESencrypt(data, key, 1, iv)
            self.tripleDESdecrypt(encrypt, key, 1, iv)
            print('Time ECB: {0} sec'.format(time.time() - start))
            start = time.time()
            encrypt = self.tripleDESencrypt(data, key, 2, iv)
            self.tripleDESdecrypt(encrypt, key, 2, iv)
            print('Time CBC: {0} sec'.format(time.time() - start))
            start = time.time()
            encrypt = self.tripleDESencrypt(data, key, 8, iv)
            self.tripleDESdecrypt(encrypt, key, 8, iv)
            print('Time innerCBC: {0} sec'.format(time.time() - start))
            start = time.time()
            encrypt = self.tripleDESencrypt(data, key, 9, iv)
            self.tripleDESdecrypt(encrypt, key, 9, iv)
            print('Time ECBpad: {0} sec'.format(time.time() - start))





if __name__ == "__main__":
    cipher = DES3Crypto()
    cipher.getMode()
    iv = Random.new().read(DES3.block_size)
    key = '\1'*16
    data = open('add.png', 'rb').read()
    cipher.testTime(data, key, iv)
    #cipher.tripleDESencrypt('input.docx', 'encrypt.txt', key, 1, iv)
    #cipher.tripleDESdecrypt('encrypt.txt', 'decrypt.docx', key, 1, iv)