Add class to encrypt the pickles

salt/crypt.py

@@ -6,25 +6,21 @@ authenticating peers
 
 # Import python libs
 import os
+import hmac
 import tempfile
 import random
 import hashlib
 import string
-# Import M2Crypto libs
+import cPickle as pickle
+# Import Cryptography libs
 from M2Crypto import RSA
+from Crypto.Cipher import AES
 # Import zeromq libs
 import zmq
 # Import salt utils
 import salt.utils
 import salt.payload
 
-def gen_aes():
-    '''
-    Generate a random aes key
-    '''
-    return hashlib.sha512(''.join(random.choice(string.ascii_letters +\
-        string.digits) for x in range(random.randint(48, 128)))).digest()
-
 class Auth(object):
     '''
     The Auth class provides the sequence for setting up communication with the
@@ -124,3 +120,77 @@ class Auth(object):
         auth['master_publish_port'] = ret['master_publish_port']
         return auth
 
+
+class AuthenticationError(Exception): pass
+
+class Crypticle(object):
+    '''
+    Authenticated encryption class
+    
+    Encryption algorithm: AES-CBC
+    Signing algorithm: HMAC-SHA256
+    '''
+
+    PICKLE_PAD = 'pickle::'
+    AES_BLOCK_SIZE = 16
+    SIG_SIZE = hashlib.sha256().digest_size
+
+    def __init__(self, key_string, key_size=192):
+        self.keys = self.extract_keys(key_string, key_size)
+        self.key_size = key_size
+
+    @classmethod
+    def generate_key_string(cls, key_size=192):
+        key = os.urandom(key_size / 8 + cls.SIG_SIZE)
+        return key.encode('base64').replace('\n', '')
+
+    @classmethod
+    def extract_keys(cls, key_string, key_size):
+        key = key_string.decode('base64')
+        assert len(key) == key_size / 8 + cls.SIG_SIZE, 'invalid key'
+        return key[:-cls.SIG_SIZE], key[-cls.SIG_SIZE:]
+
+    def encrypt(self, data):
+        '''
+        encrypt data with AES-CBC and sign it with HMAC-SHA256
+        '''
+        aes_key, hmac_key = self.keys
+        pad = self.AES_BLOCK_SIZE - len(data) % self.AES_BLOCK_SIZE
+        data = data + pad * chr(pad)
+        iv_bytes = os.urandom(self.AES_BLOCK_SIZE)
+        cypher = AES.new(aes_key, AES.MODE_CBC, iv_bytes)
+        data = iv_bytes + cypher.encrypt(data)
+        sig = hmac.new(hmac_key, data, hashlib.sha256).digest()
+        return data + sig
+
+    def decrypt(self, data):
+        '''
+        verify HMAC-SHA256 signature and decrypt data with AES-CBC
+        '''
+        aes_key, hmac_key = self.keys
+        sig = data[-self.SIG_SIZE:]
+        data = data[:-self.SIG_SIZE]
+        if hmac.new(hmac_key, data, hashlib.sha256).digest() != sig:
+            raise AuthenticationError('message authentication failed')
+        iv_bytes = data[:self.AES_BLOCK_SIZE]
+        data = data[self.AES_BLOCK_SIZE:]
+        cypher = AES.new(aes_key, AES.MODE_CBC, iv_bytes)
+        data = cypher.decrypt(data)
+        return data[:-ord(data[-1])]
+
+    def dumps(self, obj, pickler=pickle):
+        '''
+        pickle and encrypt a python object
+        '''
+        return self.encrypt(self.PICKLE_PAD + pickler.dumps(obj))
+
+    def loads(self, data, pickler=pickle):
+        '''
+        decrypt and unpickle a python object
+        '''
+        data = self.decrypt(data)
+        # simple integrity check to verify that we got meaningful data
+        assert data.startswith(self.PICKLE_PAD), 'unexpected header'
+        return pickler.loads(data[len(self.PICKLE_PAD):])
+
+