pmacs3/aes.py

import Crypto.Hash.SHA256
import Crypto.Cipher.AES

class Crypter(object):
    ALIGNMENT = 16
    def __init__(self, password, salt='aes.py'):
        self.password = password
        self.salt     = salt
        self.hash     = Crypto.Hash.SHA256.new(password + salt)
        self.cipher   = Crypto.Cipher.AES.new(self.hash.digest())

    def pad(self, s):
        xtra = len(s) % self.ALIGNMENT
        if xtra:
            return s + '\x00' * (self.ALIGNMENT - xtra)
        else:
            return s
    
    def unpad(self, s):
        l = len(s)
        while l > 0 and s[l - 1] == '\x00':
            l -= 1
        return s[:l]
    
    def encrypt(self, data):
        return self.cipher.encrypt(self.pad(data))
    
    def decrypt(self, data):
        return self.unpad(self.cipher.decrypt(self.pad(data)))