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
        self.hash      = Crypto.Hash.SHA256.new(password + salt)
        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 + '\x00' * (self.alignment - xtra)
        else:
            return s
    
    def unpad(self, s):
        '''Remove NULL padding to restore original data'''
        l = len(s)
        while l > 0 and s[l - 1] == '\x00':
            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)))