__version__ = "0.2-2"

import base64

from Crypto.Cipher import AES, ARC2, Blowfish, CAST, DES, DES3
from Crypto.Random import get_random_bytes

class Cryptmd1Error(Exception):
    pass

def add_byte_padding(content, block_size):
    l = len(content)
    d = block_size - (l % block_size)
    suffix = chr(0) * (d - 1) + chr(d)
    return content + suffix

def remove_byte_padding(content):
    c = ord(content[-1])
    return content[:-c]

add_padding = add_byte_padding

remove_padding = remove_byte_padding

ciphers = {
    "aes": {"key_sizes": (16, 24, 32), "class": AES},
    "arc2": {"key_sizes": None, "class": ARC2},
    "blowfish": {"key_sizes": None, "class": Blowfish},
    "cast": {"key_sizes": None, "class": CAST},
    "des": {"key_sizes": (8,), "class": DES},
    "des3": {"key_sizes": (16,), "class": DES3},
    }

def key_generate(size, as_ascii=False):
    key = get_random_bytes(size)
    return key

def key_generate_ascii(size):
    key = key_generate(size)
    return base64.b64encode(key)

def key_from_ascii(ascii_key):
    return base64.b64decode(ascii_key)

def key_from_ascii_file(filename):
    with open(filename, "r") as fil:
        data = fil.read()
        key = key_from_ascii(data)
    return key

def xor_blocks(block1, block2):
    return "".join((chr(ord(x) ^ ord(y))) for (x, y) in zip(block1, block2))

def cbc_encrypt(file_in, file_out, key, cipher_class=AES):
    cipher = cipher_class.new(key)
    block_size = cipher.block_size
    ve = get_random_bytes(block_size)
    file_out.write(ve)

    input_end = False
    while not input_end:
        data = file_in.read(block_size)

        if len(data) < block_size:
            input_end = True
            data = add_padding(data, block_size)

        data = xor_blocks(data, ve)
        e = cipher.encrypt(data)
        file_out.write(e)
        ve = e

def cbc_decrypt(file_in, file_out, key, cipher_class=AES):
    cipher = cipher_class.new(key)
    block_size = cipher.block_size
    ve = file_in.read(block_size)

    d = ""
    while True:
        data = file_in.read(block_size)
        l = len(data)
        assert l == block_size or l == 0
        if l == 0:
            file_out.write(remove_padding(d))
            break
        else:
            file_out.write(d)

        d = cipher.decrypt(data)
        d = xor_blocks(d, ve)
        ve = data

def steal_enc(cipher, pn1, pn, cn2):
    m = len(pn)
    b = len(pn1)

    xn1 = xor_blocks(pn1, cn2)
    en1 = cipher.encrypt(xn1)
    cn = en1[:m]
    p = pn + chr(0) * (b - m)
    dn = xor_blocks(en1, p)
    cn1 = cipher.encrypt(dn)

    return cn1, cn

def steal_dec(cipher, cn1, cn, cn2):
    m = len(cn)
    b = len(cn1)

    dn = cipher.decrypt(cn1)
    c = cn + chr(0) * (b - m)
    xn = xor_blocks(dn, c)
    pn = xn[:m]
    en1 = cn + xn[-(b - m):]
    xn1 = cipher.decrypt(en1)
    pn1 = xor_blocks(xn1, cn2)

    return pn1, pn

def cbc_steal_encrypt(file_in, file_out, key, cipher_class=AES):
    cipher = cipher_class.new(key)
    block_size = cipher.block_size
    ve = chr(0) * block_size

    d0 = file_in.read(block_size)
    if len(d0) != block_size:
        raise Cryptmd1Error("Input must be at least %d bytes" % block_size)

    while True:
        d1 = file_in.read(block_size)
        if len(d1) < block_size:
            cn1, cn = steal_enc(cipher, d0, d1, ve)
            file_out.write(cn1 + cn)
            return

        d = xor_blocks(d0, ve)
        e = cipher.encrypt(d)
        file_out.write(e)
        ve = e

        d0 = d1

def cbc_steal_decrypt(file_in, file_out, key, cipher_class=AES):
    cipher = cipher_class.new(key)
    block_size = cipher.block_size
    ve = chr(0) * block_size

    d0 = file_in.read(block_size)
    if len(d0) != block_size:
        raise Cryptmd1Error("Input must be at least %d bytes" % block_size)

    while True:
        d1 = file_in.read(block_size)
        if len(d1) < block_size:
            pn1, pn = steal_dec(cipher, d0, d1, ve)
            file_out.write(pn1 + pn)
            return

        d = cipher.decrypt(d0)
        d = xor_blocks(d, ve)
        file_out.write(d)

        ve = d0
        d0 = d1

def cfb_encrypt(file_in, file_out, key, cipher_class=AES):
    cipher = cipher_class.new(key)
    block_size = cipher.block_size

    ve = get_random_bytes(block_size)
    file_out.write(ve)

    data = file_in.read(block_size)
    while len(data) > 0:
        enc = cipher.encrypt(ve)
        d = xor_blocks(enc, data)
        file_out.write(d)
        ve = d

        data = file_in.read(block_size)

def cfb_decrypt(file_in, file_out, key, cipher_class=AES):
    cipher = cipher_class.new(key)
    block_size = cipher.block_size
    ve = file_in.read(block_size)

    data = file_in.read(block_size)
    while len(data) > 0:
        dec = cipher.encrypt(ve)
        d = xor_blocks(dec, data)
        file_out.write(d)

        ve = data
        data = file_in.read(block_size)

crypt_functions = {
    "cbc": {"encrypt": cbc_encrypt, "decrypt": cbc_decrypt},
    "cbc-steal": {"encrypt": cbc_steal_encrypt, "decrypt": cbc_steal_decrypt},
    "cfb": {"encrypt": cfb_encrypt, "decrypt": cfb_decrypt},
    }

def cipher_class_from_str(key):
    return ciphers[key]["class"]