rootninja

I needed to encrypt some strings before sending them out over the network to another process and I don’t want to rely on another mechanism for handling the encryption, so here’s my simple AES encryption class. You can load plain text out of the vault by providing a cipher or save an encrypted message by supplying the plain text.

Using AES.MODE_CBC (Cipher-Block Chaining mode) each block of text is XOR’d with the previous cipher block and then encrypted. The result being each block is dependent on the plain text up until that point. This requires specific predetermined lengths however…

#!/usr/bin/env python
from Crypto.Cipher import AES
import hashlib
import os

class Vault():
    def __init__(self, key):
        '''AES encryption and decryption vault'''
        mode = AES.MODE_CBC
        key = self.set_key_length(key)
        key = hashlib.sha256(key).digest()
        self.encryptor = AES.new(key, mode)
        self.decryptor = AES.new(key, mode)

    def encr(self, text):
        '''encrypt plain text, return the cipher'''
        plain = self.set_text_length(text)
        cipher = self.encryptor.encrypt(plain)
        return cipher

    def decr(self, cipher):
        '''decrpyt a cipher, return plain text'''
        plain = self.decryptor.decrypt(cipher)
        plain = plain.rstrip()
        return plain

    def set_text_length(self, text):
        '''Force text length into a multiple of 16'''
        while (len(text)%16) is not 0:
            text += ' '
        return text

    def set_key_length(self, key):
        '''Force key length to 16, 24, or 32'''
        if len(key) > 32:
            key = key[:32]
        else:
            while (len(key)%32) is not 0:
                key += ' '
        return key

    def save(self, text, filename):
        try:
            f = open(os.path.join(os.path.dirname(__file__), filename), 'w')
            cipher = self.encr(text)
            f.write(cipher)
            f.close()
            return cipher
        except:
            return False

    def load(self, cipher, filename):
        try:
            f = open(os.path.join(os.path.dirname(__file__), filename), 'r')
            text = self.decr(f.read())
            f.close()
            return text
        except:
            return False

if __name__ == "__main__":
    string_list = [ 'test123', '9', 'asdf jkl;', '!@#$)(%^' ]
    key = 'super_secret_encryption_key_string'
    v = Vault(key)

    for text in string_list:
        print 'text/key = "%s" "%s"' % (text, key)
        cipher = v.encr(text)
        print 'Encrypted text = "%s"' % cipher
        plain = v.decr(cipher)
        print 'Decrypted text =  "%s"' % plain
        print ''

    text = 'plain-and-cipher-testing-save-and-load-from-file'
    print 'Plain text = "%s"' % text

    key = 'yellow'
    print 'Key = "%s"' % key

    filename = 'password.enc'
    print 'Saving to filename: "%s"' % filename

    saved_cipher = v.save(text, filename)
    print 'saved cipher = "%s"' % saved_cipher

    loaded_text = v.load(saved_cipher, filename)
    print 'loaded text = "%s"' % loaded_text

If you change the mode to AES.MODE_CFB (Cipher Feedback mode) you can remove the length restrictions on the key and text completely. In this mode the text is encrypted one character at a time instead of by block. No need to check and adjust the key size before calling AES.new(key, mode) and for the plain text, just can encrypt it directly without fixing its length first.