AES Class Python

aes.py

#!/usr/bin/python2.7
# -*- coding: utf-8 -*-

import string
import random
import base64
from Crypto.Cipher import AES

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


class AESCipher:

    def __init__(self, KEY=u'0123456789abcdef0123456789abcdef', IV='12345678abcdefgh'):
        self.key = KEY
        self.iv = IV.zfill(16)
        #self.mode = AES.MODE_ECB
        self.mode = AES.MODE_CBC

    # ----------
    def encrypt(self, raw, base64encode=False):
        cipher = AES.new(self.key, self.mode, self.iv)
        encode = cipher.encrypt(pad(raw))
        if base64encode:
            return base64.b64encode(encode)
        return encode

    def decrypt(self, enc, base64encode=False):
        cipher = AES.new(self.key, self.mode, self.iv)
        if base64encode:
            try:
                return unpad(cipher.decrypt(base64.b64decode(enc)))
            except Exception:
                return None
        return unpad(cipher.decrypt(enc))

    # ----------
    def encrypt_random(self, raw, base64encode=False):
        iv = ''.join(random.choice(string.ascii_uppercase) for i in range( AES.block_size ))
        cipher = AES.new( self.key, self.mode, iv )
        encode = iv + cipher.encrypt( pad(raw) )
        if base64encode:
            return base64.b64encode(encode)
        return encode

    def decrypt_random( self, enc, base64encode=False):
        if base64encode:
            try:
                enc = base64.b64decode(enc)
            except Exception:
                return None
        cipher = AES.new(self.key, self.mode, enc[:16] )
        return unpad( cipher.decrypt(enc[16:]) )