move key loading to new class Keys

README.md

@@ -83,6 +83,12 @@ Disadvantages
 Other encrypted field modules are available if you just want to use encrypted field classes in Django models and do not need unique constraints nor plan to join tables on encrypted fields for analysis.
 
 
+Troubleshooting
+---------------
+
+If you are adding _django-crypto_fields_ to an existing project, model loading can get in the way if you already added encrypted fields to a model. To get around this, try renaming the models.py (or module folder) before generating the keys. Once the keys are generated you can name the module back to models.   
+
+
 Contribute
 ----------
 

django_crypto_fields/__init__.py

@@ -0,0 +1 @@
+default_app_config = 'django_crypto_fields.apps.DjangoCryptoFieldsConfig'

django_crypto_fields/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class DjangoCryptoFieldsConfig(AppConfig):
+    name = 'django_crypto_fields'
+    verbose_name = "django-crypto-fields"

django_crypto_fields/classes/cryptor.py

@@ -1,21 +1,12 @@
-import copy
-import logging
+import sys
 
 from Crypto import Random
-from Crypto.PublicKey import RSA
-from Crypto.Cipher import PKCS1_OAEP, AES
-from Crypto.Util import number
+from Crypto.Cipher import AES
 
 from ..constants import KEY_FILENAMES, ENCODING
 from ..exceptions import EncryptionError
 
-logger = logging.getLogger(__name__)
-
-
-class NullHandler(logging.Handler):
-    def emit(self, record):
-        pass
-nullhandler = logger.addHandler(NullHandler())
+from .keys import KEYS
 
 
 class Cryptor(object):
@@ -24,31 +15,28 @@ class Cryptor(object):
     The PEM file names and paths are in KEY_FILENAMES. KEYS is a copy of this except the
     filenames are replaced with the actual keys."""
 
-    KEYS = copy.deepcopy(KEY_FILENAMES)
-
     def __init__(self):
-        self.rsa_key_info = {}
-        self.load_keys()
+        pass
 
     def aes_encrypt(self, plaintext, mode):
         try:
             plaintext = plaintext.encode(ENCODING)
         except AttributeError:
             pass
-        aes_key = self.KEYS.get('aes').get(mode).get('private')
+        aes_key = KEYS['aes'][mode]['private']
         iv = Random.new().read(AES.block_size)
         cipher = AES.new(aes_key, AES.MODE_CFB, iv)
         return iv + cipher.encrypt(plaintext)
 
     def aes_decrypt(self, ciphertext, mode):
-        aes_key = self.KEYS.get('aes').get(mode).get('private')
+        aes_key = KEYS['aes'][mode]['private']
         iv = ciphertext[:AES.block_size]
         cipher = AES.new(aes_key, AES.MODE_CFB, iv)
         plaintext = cipher.decrypt(ciphertext)[AES.block_size:]
         return plaintext.decode(ENCODING)
 
     def rsa_encrypt(self, plaintext, mode):
-        rsa_key = self.KEYS.get('rsa').get(mode).get('public')
+        rsa_key = KEYS['rsa'][mode]['public']
         try:
             plaintext = plaintext.encode(ENCODING)
         except AttributeError:
@@ -60,78 +48,33 @@ def rsa_encrypt(self, plaintext, mode):
         return ciphertext
 
     def rsa_decrypt(self, ciphertext, mode):
-        rsa_key = self.KEYS.get('rsa').get(mode).get('private')
+        rsa_key = KEYS['rsa'][mode]['private']
         plaintext = rsa_key.decrypt(ciphertext)
         return plaintext.decode(ENCODING)
 
-    def update_rsa_key_info(self, rsa_key, mode):
-        """Stores info about the RSA key."""
-        modBits = number.size(rsa_key._key.n)
-        self.rsa_key_info[mode] = {'bits': modBits}
-        k = number.ceil_div(modBits, 8)
-        self.rsa_key_info[mode].update({'bytes': k})
-        hLen = rsa_key._hashObj.digest_size
-        self.rsa_key_info[mode].update({'max_message_length': k - (2 * hLen) - 2})
-
-    def load_keys(self):
-        logger.info('/* Loading keys ...')
-        try:
-            # load RSA
-            for mode, keys in KEY_FILENAMES['rsa'].items():
-                for key in keys:
-                    key_file = KEY_FILENAMES['rsa'][mode][key]
-                    with open(key_file, 'rb') as f:
-                        rsa_key = RSA.importKey(f.read())
-                        rsa_key = PKCS1_OAEP.new(rsa_key)
-                        self.KEYS['rsa'][mode][key] = rsa_key
-                        self.update_rsa_key_info(rsa_key, mode)
-                    logger.info('(*) Loaded ' + key_file)
-            # decrypt and load AES
-            for mode in KEY_FILENAMES['aes']:
-                rsa_key = self.KEYS['rsa'][mode]['private']
-                key_file = KEY_FILENAMES['aes'][mode]['private']
-                with open(key_file, 'rb') as faes:
-                    aes_key = rsa_key.decrypt(faes.read())
-                self.KEYS['aes'][mode]['private'] = aes_key
-                logger.info('(*) Loaded ' + key_file)
-            # decrypt and load salt
-            for mode in KEY_FILENAMES['salt']:
-                rsa_key = self.KEYS['rsa'][mode]['private']
-                key_file = KEY_FILENAMES['salt'][mode]['private']
-                with open(key_file, 'rb') as fsalt:
-                    salt = rsa_key.decrypt(fsalt.read())
-                self.KEYS['salt'][mode]['private'] = salt
-                logger.info('(*) Loaded ' + key_file)
-            logger.info('Done preloading keys. */')
-        except FileNotFoundError:
-            raise FileNotFoundError(
-                'Unable to find keys. Check the KEY_PATH or, '
-                'if you do not have any keys, run the \'generate_keys\' management command '
-                'and try again.')
-
     def test_rsa(self):
         """ Tests keys roundtrip"""
         plaintext = 'erik is a pleeb!'
-        for mode in KEY_FILENAMES.get('rsa'):
+        for mode in KEY_FILENAMES['rsa']:
             try:
-                rsa_key = self.KEYS.get('rsa').get(mode).get('public')
+                rsa_key = KEYS['rsa'][mode]['public']
                 ciphertext = rsa_key.encrypt(plaintext.encode('utf_8'))
-                print('(*) Passed encrypt: ' + KEY_FILENAMES.get('rsa').get(mode).get('public'))
+                sys.stdout.write('(*) Passed encrypt: ' + KEY_FILENAMES['rsa'][mode]['public'])
             except (AttributeError, TypeError) as e:
                 print('( ) Failed encrypt: {} public ({})'.format(mode, e))
             try:
-                rsa_key = self.KEYS.get('rsa').get(mode).get('private')
+                rsa_key = KEYS['rsa'][mode]['private']
                 assert plaintext == rsa_key.decrypt(ciphertext).decode(ENCODING)
-                print('(*) Passed decrypt: ' + KEY_FILENAMES.get('rsa').get(mode).get('private'))
+                print('(*) Passed decrypt: ' + KEY_FILENAMES['rsa'][mode]['private'])
             except (AttributeError, TypeError) as e:
                 print('( ) Failed decrypt: {} private ({})'.format(mode, e))
 
     def test_aes(self):
         """ Tests keys roundtrip"""
         plaintext = 'erik is a pleeb!'
-        for mode in KEY_FILENAMES.get('aes'):
+        for mode in KEY_FILENAMES['aes']:
             ciphertext = self.aes_encrypt(plaintext, mode)
             assert plaintext != ciphertext
-            print('(*) Passed encrypt: ' + KEY_FILENAMES.get('aes').get(mode).get('private'))
+            print('(*) Passed encrypt: ' + KEY_FILENAMES['aes'][mode]['private'])
             assert plaintext == self.aes_decrypt(ciphertext, mode)
-            print('(*) Passed decrypt: ' + KEY_FILENAMES.get('aes').get(mode).get('private'))
+            print('(*) Passed decrypt: ' + KEY_FILENAMES['aes'][mode]['private'])

django_crypto_fields/classes/field_cryptor.py

@@ -9,6 +9,7 @@
 from ..exceptions import CipherError, EncryptionError, MalformedCiphertextError, EncryptionKeyError
 
 from .cryptor import Cryptor
+from .keys import KEYS
 
 
 class FieldCryptor(object):
@@ -46,7 +47,7 @@ def hash(self, plaintext):
         except AttributeError:
             pass
         try:
-            salt = self.cryptor.KEYS.get('salt').get(self.mode).get('private')
+            salt = KEYS['salt'][self.mode]['private']
         except AttributeError:
             raise EncryptionKeyError('Invalid mode for salt key. Got {}'.format(self.mode))
         dk = hashlib.pbkdf2_hmac(HASH_ALGORITHM, plaintext, salt, HASH_ROUNDS)

django_crypto_fields/classes/keys.py

@@ -0,0 +1,78 @@
+import copy
+import sys
+
+from Crypto.Cipher import PKCS1_OAEP
+from Crypto.PublicKey import RSA
+from Crypto.Util import number
+
+from ..constants import KEY_FILENAMES, KEY_PATH, KEY_PREFIX
+from ..utils import KeyGenerator
+
+KEYS = copy.deepcopy(KEY_FILENAMES)
+
+
+class Keys(object):
+
+    def __init__(self):
+        self.loaded = False
+        self.rsa_key_info = {}
+        try:
+            self.load_keys()
+        except FileNotFoundError:
+            sys.stdout.write('/* Loading keys failed.\n')
+            KeyGenerator.create_keys(KEY_PATH, KEY_PREFIX, show_msgs=False)
+            self.load_keys()
+
+    def load_rsa_key(self, mode, key):
+        """Loads an RSA key."""
+        key_file = KEY_FILENAMES['rsa'][mode][key]
+        with open(key_file, 'rb') as frsa:
+            rsa_key = RSA.importKey(frsa.read())
+            rsa_key = PKCS1_OAEP.new(rsa_key)
+            KEYS['rsa'][mode][key] = rsa_key
+            self.update_rsa_key_info(rsa_key, mode)
+        return key_file
+
+    def load_aes_key(self, mode, key):
+        """Decrypts and loads an AES key."""
+        rsa_key = KEYS['rsa'][mode]['private']
+        key_file = KEY_FILENAMES['aes'][mode]['private']
+        with open(key_file, 'rb') as faes:
+            aes_key = rsa_key.decrypt(faes.read())
+        KEYS['aes'][mode]['private'] = aes_key
+        return key_file
+
+    def load_salt_key(self, mode, key):
+        """Decrypts and loads a salt key."""
+        rsa_key = KEYS['rsa'][mode]['private']
+        key_file = KEY_FILENAMES['salt'][mode]['private']
+        with open(key_file, 'rb') as fsalt:
+            salt = rsa_key.decrypt(fsalt.read())
+        KEYS['salt'][mode]['private'] = salt
+        return key_file
+
+    def load_keys(self):
+        sys.stdout.write('/* Loading keys ...\n')
+        for mode, keys in KEY_FILENAMES['rsa'].items():
+            for key in keys:
+                key_file = self.load_rsa_key(mode, key)
+                sys.stdout.write('(*) Loaded {}\n'.format(key_file))
+        for mode in KEY_FILENAMES['aes']:
+            key_file = self.load_aes_key(mode, key)
+            sys.stdout.write('(*) Loaded {}\n'.format(key_file))
+        for mode in KEY_FILENAMES['salt']:
+            key_file = self.load_salt_key(mode, key)
+            sys.stdout.write('(*) Loaded {}\n'.format(key_file))
+        sys.stdout.write('Done loading keys.\n')
+        self.loaded = True
+
+    def update_rsa_key_info(self, rsa_key, mode):
+        """Stores info about the RSA key."""
+        modBits = number.size(rsa_key._key.n)
+        self.rsa_key_info[mode] = {'bits': modBits}
+        k = number.ceil_div(modBits, 8)
+        self.rsa_key_info[mode].update({'bytes': k})
+        hLen = rsa_key._hashObj.digest_size
+        self.rsa_key_info[mode].update({'max_message_length': k - (2 * hLen) - 2})
+
+keys = Keys()

django_crypto_fields/fields/base_field.py

@@ -2,6 +2,7 @@
 from django.db import models
 
 from ..classes import FieldCryptor
+from ..classes.keys import keys
 from ..constants import HASH_PREFIX, ENCODING
 from ..exceptions import CipherError, EncryptionError, MalformedCiphertextError
 
@@ -16,7 +17,7 @@ def __init__(self, *args, **kwargs):
         self.field_cryptor = FieldCryptor(algorithm, mode)
         max_length = kwargs.get('max_length', None) or len(HASH_PREFIX) + self.field_cryptor.hash_size
         if algorithm == 'rsa':
-            max_message_length = self.field_cryptor.cryptor.rsa_key_info[mode]['max_message_length']
+            max_message_length = keys.rsa_key_info[mode]['max_message_length']
             if max_length > max_message_length:
                 raise EncryptionError(
                     '{} attribute \'max_length\' cannot exceed {} for RSA. Got {}. '