import hashlib,os
from utility import read_in_chunks


try:
    pycrypto = True
    from Crypto.Cipher import AES
    from Crypto import Random
    print("Using PyCrypto for AES enc/dec")
except ImportError:
    pycrypto = False
    import pyaes
    print("Using Pyaes for AES enc/dec")



class AESCipherFile:
    def __init__( self ,f,d, size_limit_reading_os, key ):
        h = hashlib.sha256()
        h.update(key.encode("utf_8"))
        self.key = h.digest()
        self.f = f
        self.d = d
        self.size_limit_reading_os = size_limit_reading_os

    def pad(self,b):
        b = bytearray(b)
        fill = 16 - len(b) % 16
        for i in range(fill):
            b.append(fill)
        return bytes(b)

    def unpad(self,b) :
        return b[:-b[len(b)-1]]


    def encrypt( self ):
        if pycrypto:
            iv = Random.new().read( 16 )
            aes_encrypt_state = AES.new( self.key, AES.MODE_CBC, iv )
        else:
            iv = os.urandom(16)
            aes_encrypt_state = pyaes.AESModeOfOperationCBC(self.key, iv = iv)
        #check if the file is a multiple of 16 bytes for padding
        self.f.seek(0,2)
        if self.f.tell() % 16 == 0:
            self.d.write(bytes(b"1"))
        else:
            self.d.write(bytes(b"0"))
        self.f.seek(0,0)
        #end check
        self.d.write(iv)
        hash = hashlib.md5()
        for piece in read_in_chunks(self.f,self.size_limit_reading_os):
            hash.update(piece)
            if len(piece) % 16 > 0 :
                piece = self.pad(piece)
            steps = int(len(piece) / 16)
            for i in range(steps):
                self.d.write(aes_encrypt_state.encrypt( piece[16 * i : 16 * (i + 1)] ))
        return hash.hexdigest()

    def decrypt( self ):
        perfect = False
        perfectbyte = self.f.read(1)
        if perfectbyte == bytes(b"1"):
            perfect = True
        if pycrypto:
            aes_decrypt_state = AES.new( self.key, AES.MODE_CBC, self.f.read(16) )
        else:
            aes_decrypt_state = pyaes.AESModeOfOperationCBC(self.key, self.f.read(16))
        last_16_bytes_in_piece = None
        for piece in read_in_chunks(self.f,self.size_limit_reading_os):
            if last_16_bytes_in_piece != None:
               self.d.write(last_16_bytes_in_piece)
            steps = int(len(piece) / 16)
            if steps > 1 :
                for i in range(steps - 1):
                    self.d.write(aes_decrypt_state.decrypt( piece[16 * i : 16 * (i + 1)] ))
                i = i + 1
            else:
                i = 0
            last_16_bytes_in_piece = aes_decrypt_state.decrypt( piece[16 * i : 16 * (i + 1)] )
        if last_16_bytes_in_piece != None:
            if perfect:
                self.d.write(last_16_bytes_in_piece)
            else:
                self.d.write(self.unpad(last_16_bytes_in_piece))



''' usage
with open("c:\\test\\file", 'rb') as f:
    with open("c:\\test\\file_enc", 'wb') as d:
        aes = AESCipherFile(f,d,134217728,"chiave")
        aes.encrypt()
with open("c:\\test\\file_enc", 'rb') as f:
    with open("c:\\test\\file", 'wb') as d:
        aes = AESCipherFile(f,d,134217728,"chiave")
        aes.decrypt()
'''