#-*- coding: utf-8 -*-
from aespython import key_expander, aes_cipher, ocb_mode
import json
import base64

class sjcl(ocb_mode.OCBMode):
    """convenience class to wrap the sjcl objects encryption and decryption"""
    def __init__(self, key, keylen):
        KE = key_expander.KeyExpander(keylen)
        expended_key = KE.expand(key)
        cipher = aes_cipher.AESCipher(expanded_key)
        ocb_mode.OCBMode.__init__(self, cipher, len(key))

    @staticmethod
    def hex2byte(hexstring):
        """convert a hexstring to a list of bytes"""
        return map(ord, hexstring.decode("hex"))

    @staticmethod
    def byte2hex(byteslist):
        """convert a list of bytes to a hexstring"""
        return str(bytearray(byteslist)).encode("hex")
    
    @staticmethod
    def pad(b64string):
        """pad a base64 string with = so that its lengh is divisible by 4"""
        return b64string + ''.join(['='] * (len(b64string) % 4))

    def encrypt_sjclobject(self, plaintext, header):
        """encrypt a message and return in in the sjcl object format"""
        tag, ciphertext = ocb_mode.OCBMode.encrypt_block(self, plaintext, header)
        sjcl_object = {"iv": base64.b64encode(bytearray(self._iv)).rstrip("="),
                       "mode": "ocb2",
                       "ts": 128,
                       "ct": base64.b64encode(bytearray(ciphertext) +
                                              bytearray(tag)).rstrip("="),
                       "adata": base64.b64encode(header).rstrip("=")
                  }
        return json.dumps(sjcl_object)

    def decrypt_sjclobject(self, objectstring):
        """takes a sjcl object and return a list of bytes."""
        sjcl_object = json.loads(objectstring)
        iv = map(ord, base64.b64decode(pad(sjcl_object['iv'])))
        self.set_iv(iv)
        ct = map(ord, base64.b64decode(pad(sjcl_object['ct'])))
        try:
            header = map(ord, base64.b64decode(pad(sjcl_object['adata'])))
        except KeyError:
            header = []
        #sjcl appends the tag at the end of the cipher text (it's 16 bytes long)
        tag = ct[-16:]
        ciphertext = ct[:-16]
        isValid, plaintext = ocb_mode.OCBMode.decrypt_block(self, header, ciphertext, tag)
        if not isValid:
            raise Exception("can't authenticate the data")
        else:
            return plaintext
        
if __name__=="__main__":
    key = sjcl.hex2byte('12538243c49f1c58e6f7b0687bbd65b2')
    iv = sjcl.hex2byte('250c3041c00a605a4100e264abbc588b')
    sjclmode = sjcl(key, 128)
    sjclmode.set_iv(iv)
    plaintext = map(ord, "La chaire est triste, hélas ! et j'ai lu \
tous les livres.")
    header = bytearray("Mallarmé")
    test = sjclmode.encrypt_sjclobject(plaintext, header)   
    print test
    print str(bytearray(sjclmode.decrypt_sjclobject(test)))