aes.py3: Simple AES helper

aes.py

"""
"""

from Crypto import Random
from Crypto.Cipher import AES

import pkcs7


random = Random.new().read


def gen_keyiv(keylen=32, ivlen=16):
    """
    """
    return (random(keylen), random(ivlen))


# @param    bstr    text
# @param    bstr    key
# @param    bstr    iv
# @return   bstr  
def decrypt(ctext, key, iv, mode=AES.MODE_CBC):
    """
    """
    cipher = AES.new(key, mode, iv) # A new cipher each time is required.
    padded_text = cipher.decrypt(ctext)
    return pkcs7.decode(padded_text)


# @param    str     text
# @param    bstr    key
# @param    bstr    iv
# @return   bstr    
def encrypt(text, key, iv, mode=AES.MODE_CBC):
    cipher = AES.new(key, mode, iv) # A new cipher each time is required.
    padded_text = pkcs7.encode(text)
    return cipher.encrypt(padded_text)


class Aes(object):
    """Simply hangs on to a key/iv and exposes `encrypt` and `decrypt`.
    """
    def __init__(self, key, iv):
        self.key, self.iv = key, iv

    def encrypt(self, text):
        return encrypt(text, self.key, self.iv)

    def decrypt(self, ctext):
        return decrypt(ctext, self.key, self.iv)


# Tests
# =====
import unittest

class AesTest(unittest.TestCase):
    def test_aes(self):
        import aes


        def test(text, key, iv):
            a = aes.Aes(key, iv)
            assert a.decrypt(a.encrypt(text)) == text


        test(b'iviviviviviviviviviviviviviviviv',
                b'iviviviviviviviviviviviviviviviv',
                b'iviviviviviviviv')

        test('01234567890'.encode(), *aes.gen_keyiv())