pmacs3/aesutil.py

try:
    import Crypto.Hash.SHA256
    import Crypto.Cipher.AES
    has_aes = True
except:
    has_aes = False

class CrypterInitError(Exception):
    pass

class Crypter(object):
    '''Wraps AES functionality provided by pycrypto'''
    def __init__(self, password, salt='aes.py', alignment=16):
        '''Initializes a crypter object'''
        if not has_aes:
            raise CrypterInitError("pycrypto not installed")
        self.password  = password
        self.salt      = salt
        pdata = (password + salt).encode('UTF-8')
        self.hash      = Crypto.Hash.SHA256.new(pdata)
        self.cipher    = Crypto.Cipher.AES.new(self.hash.digest())
        self.alignment = alignment

    def pad(self, s):
        '''Add NULL padding to create aligned data'''
        xtra = len(s) % self.alignment
        if xtra:
            return s + b'\x00' * (self.alignment - xtra) #FIX3
        else:
            return s

    def unpad(self, s):
        '''Remove NULL padding to restore original data'''
        l = len(s)
        assert type(s) == type(b'')
        while l > 0 and s[l - 1] == 0: #FIX3
            l -= 1
        return s[:l]

    def encrypt(self, data):
        '''Creates encrypted data string from original input'''
        return self.cipher.encrypt(self.pad(data))

    def decrypt(self, data):
        '''Creates original data string from encrypted input'''
        return self.unpad(self.cipher.decrypt(self.pad(data)))