#!/usr/bin/env python

import base64
import sys
import os
import random 
import struct

from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto import Random

from cexceptions import *

chunk_length = 16
pad = lambda s: s + (chunk_length - len(s) % chunk_length) * chr(chunk_length - len(s) % chunk_length) 
unpad = lambda s : s[0:-ord(s[-1])]

class DeltaCrypt(object):
    key = ''
   
    def __init__(self, key):
        self.key = key.encode()
        key_length = len(self.key)
        if key_length not in [16, 24, 32]:
            raise CryptoKeyLengthException('AES key_length must be in [16, 24, 32], please verify your key: %s (%s)' % (key, key_length))
        
    def _spawn_cryptor(self, iv = ''):
        if iv:
            return AES.new(self.key, AES.MODE_CBC, iv)
        else: 
            iv = Random.new().read(AES.block_size)
            return [iv, AES.new(self.key, AES.MODE_CBC, iv)]

    def encrypt(self, raw_string, return_checksum=True):
        padded = pad(raw_string)

        iv, cryptor = self._spawn_cryptor()
    
        encrypted = base64.b64encode(iv + cryptor.encrypt(padded))

        if return_checksum:
            return [SHA256.new(raw_string).hexdigest(), encrypted]
        else:
            return encrypted
            
    def decrypt(self, encrypted_string, return_checksum=True):
        encrypted_string = base64.b64decode(encrypted_string)
        iv = encrypted_string[:16]
        
        decrypted = unpad(self._spawn_cryptor(iv).decrypt(encrypted_string[16:]))
        
        if return_checksum:
            return [SHA256.new(decrypted).hexdigest(), decrypted]
        else:
            return decrypted