from hashlib import md5

from Crypto import Random
from Crypto.Cipher import AES


class ArchiveEnc:
    def __init__(self, dir_name, dir_path, zip_path, enc_path):
        self.dir_name = dir_name
        self.dir_path = dir_path
        self.zip_path = zip_path
        self.enc_path = enc_path


def __derive_key_and_iv(password, salt, key_length, iv_length):
    d = d_i = ''
    while len(d) < key_length + iv_length:
        d_i = md5(d_i + password + salt).digest()
        d += d_i
    return d[:key_length], d[key_length:key_length + iv_length]


def encrypt_files(archive_zip_list, password):
    result_list = list()
    for archive_zip in archive_zip_list:
        print "Encrypting [{}]".format(archive_zip.zip_path)
        file_path_enc = archive_zip.zip_path + '_enc'
        with open(archive_zip.zip_path, 'rb') as in_file, open(file_path_enc, 'wb') as out_file:
            __encrypt(in_file, out_file, password)
            result_list.append(ArchiveEnc(archive_zip.dir_name, archive_zip.dir_path, archive_zip.zip_path, file_path_enc))
    return result_list


def decrypt_file(file_path, password):
    print "Decrypting [{}]".format(file_path)
    file_path_dec = file_path + '_dec'
    with open(file_path, 'rb') as in_file, open(file_path_dec, 'wb') as out_file:
        __decrypt(in_file, out_file, password)
    return file_path_dec


def __encrypt(in_file, out_file, password, key_length=32):
    bs = AES.block_size
    salt = Random.new().read(bs - len('Salted__'))
    key, iv = __derive_key_and_iv(password, salt, key_length, bs)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    out_file.write('Salted__' + salt)
    finished = False
    while not finished:
        chunk = in_file.read(1024 * bs)
        if len(chunk) == 0 or len(chunk) % bs != 0:
            padding_length = (bs - len(chunk) % bs) or bs
            chunk += padding_length * chr(padding_length)
            finished = True
        out_file.write(cipher.encrypt(chunk))


def __decrypt(in_file, out_file, password, key_length=32):
    bs = AES.block_size
    salt = in_file.read(bs)[len('Salted__'):]
    key, iv = __derive_key_and_iv(password, salt, key_length, bs)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    next_chunk = ''
    finished = False
    while not finished:
        chunk, next_chunk = next_chunk, cipher.decrypt(in_file.read(1024 * bs))
        if len(next_chunk) == 0:
            padding_length = ord(chunk[-1])
            chunk = chunk[:-padding_length]
            finished = True
        out_file.write(chunk)
