my wrapper for the AES implementation in PyCrypto

aes.py

#!/usr/bin/env python

# to get the same interfaces as https://github.com/bozhu/AES-Python

from Crypto.Cipher import AES as _AES
from Crypto.Util.number import long_to_bytes, bytes_to_long


class AES:
    def __init__(self, master_key, block_size=128):
        assert block_size in (128, 192, 256)
        self.key_size = block_size / 8
        self.change_key(master_key)

    def change_key(self, master_key):
        key_bytestr = long_to_bytes(master_key, self.key_size)
        assert len(key_bytestr) == self.key_size
        self.__aes = _AES.new(key_bytestr, _AES.MODE_ECB)

    def encrypt(self, plaintext):
        plain_bytestr = long_to_bytes(plaintext, 16)
        assert len(plain_bytestr) == 16
        cipher_bytestr = self.__aes.encrypt(plain_bytestr)
        # assert len(cipher_bytestr) == 16
        return bytes_to_long(cipher_bytestr)

    def decrypt(self, ciphertext):
        cipher_bytestr = long_to_bytes(ciphertext, 16)
        assert len(cipher_bytestr) == 16
        plain_bytestr = self.__aes.decrypt(cipher_bytestr)
        # assert len(plain_bytestr) == 16
        return bytes_to_long(plain_bytestr)


if __name__ == '__main__':
    plaintext = 0x3243f6a8885a308d313198a2e0370734
    master_key = 0x2b7e151628aed2a6abf7158809cf4f3c
    # the ciphertext should be
    # 0x3925841d02dc09fbdc118597196a0b32

    my_AES = AES(master_key)

    encrypted = my_AES.encrypt(plaintext)
    decrypted = my_AES.decrypt(encrypted)

    print 'plaintext:', hex(plaintext)
    print 'masterkey:', hex(master_key)

    print 'encrypted:', hex(encrypted),
    if encrypted == 0x3925841d02dc09fbdc118597196a0b32:
        print 'correct!'
    else:
        print 'wrong...'
    print 'should be:', hex(0x3925841d02dc09fbdc118597196a0b32)

    print 'decrypted:', hex(decrypted),
    if decrypted == plaintext:
        print 'correct!'
    else:
        print 'wrong...'