Merge pull request #32 from defrex/master

Allow Django Admin to initialize blank number fields, fix multiple object failure (with tests), and add south introspection rules for new fields introduced in previous release. (includes pull requests #33 and #34).

src/django_fields/fields.py

@@ -14,8 +14,11 @@
 from Crypto.Random import random
 
 if hasattr(settings, 'USE_CPICKLE'):
-    warnings.warn("The USE_CPICKLE options is now obsolete. cPickle will always "
-    "be used unless it cannot be found or DEBUG=True",DeprecationWarning)
+    warnings.warn(
+        "The USE_CPICKLE options is now obsolete. cPickle will always "
+        "be used unless it cannot be found or DEBUG=True",
+        DeprecationWarning,
+    )
 
 if settings.DEBUG:
     import pickle
@@ -32,8 +35,13 @@ class BaseEncryptedField(models.Field):
     def __init__(self, *args, **kwargs):
         self.cipher_type = kwargs.pop('cipher', 'AES')
         self.block_type = kwargs.pop('block_type', None)
-        if self.block_type == None:
-            warnings.warn("Default usage of pycrypto's AES block type defaults has been deprecated and will be removed in 0.3.0 (default will become MODE_CBC). Please specify a secure block_type, such as CBC.", DeprecationWarning)
+        if self.block_type is None:
+            warnings.warn(
+                "Default usage of pycrypto's AES block type defaults has been "
+                "deprecated and will be removed in 0.3.0 (default will become "
+                "MODE_CBC). Please specify a secure block_type, such as CBC.",
+                DeprecationWarning,
+            )
         try:
             imp = __import__('Crypto.Cipher', globals(), locals(), [self.cipher_type], -1)
         except:
@@ -71,7 +79,6 @@ def _get_padding(self, value):
         mod = (len(value) + 2) % self.cipher.block_size
         return self.cipher.block_size - mod + 2
 
-
     def to_python(self, value):
         if self._is_encrypted(value):
             if self.block_type:
@@ -84,9 +91,7 @@ def to_python(self, value):
             else:
                 decrypt_value = binascii.a2b_hex(value[len(self.prefix):])
             return force_unicode(
-                self.cipher.decrypt(
-                    decrypt_value
-                ).split('\0')[0]
+                self.cipher.decrypt(decrypt_value).split('\0')[0]
             )
         return value
 
@@ -97,11 +102,17 @@ def get_db_prep_value(self, value, connection=None, prepared=False):
         value = smart_str(value)
 
         if not self._is_encrypted(value):
-            padding  = self._get_padding(value)
+            padding = self._get_padding(value)
             if padding > 0:
-                value += "\0" + ''.join([random.choice(string.printable)
-                    for index in range(padding-1)])
+                value += "\0" + ''.join([
+                    random.choice(string.printable)
+                    for index in range(padding-1)
+                ])
             if self.block_type:
+                self.cipher = self.cipher_object.new(
+                    settings.SECRET_KEY[:32],
+                    getattr(self.cipher_object, self.block_type),
+                    self.iv)
                 value = self.prefix + binascii.b2a_hex(self.iv + self.cipher.encrypt(value))
             else:
                 value = self.prefix + binascii.b2a_hex(self.cipher.encrypt(value))
@@ -134,8 +145,10 @@ def formfield(self, **kwargs):
     def get_db_prep_value(self, value, connection=None, prepared=False):
         if value is not None and not self._is_encrypted(value):
             if len(value) > self.unencrypted_length:
-                raise ValueError("Field value longer than max allowed: " +
-                    str(len(value)) + " > " + str(self.unencrypted_length))
+                raise ValueError(
+                    "Field value longer than max allowed: " +
+                    str(len(value)) + " > " + str(self.unencrypted_length)
+                )
         return super(EncryptedCharField, self).get_db_prep_value(
             value,
             connection=connection,
@@ -157,7 +170,7 @@ def get_internal_type(self):
         return 'CharField'
 
     def formfield(self, **kwargs):
-        defaults = {'widget': self.form_widget,'form_class':self.form_field}
+        defaults = {'widget': self.form_widget, 'form_class': self.form_field}
         defaults.update(kwargs)
         return super(BaseEncryptedDateField, self).formfield(**defaults)
 
@@ -220,7 +233,7 @@ def get_internal_type(self):
 
     def to_python(self, value):
         # value is either an int or a string of an integer
-        if isinstance(value, self.number_type):
+        if isinstance(value, self.number_type) or value == '':
             number = value
         else:
             number_text = super(BaseEncryptedNumberField, self).to_python(value)
@@ -301,10 +314,10 @@ def formfield(self, **kwargs):
 
 class EncryptedUSSocialSecurityNumberField(BaseEncryptedField):
     __metaclass__ = models.SubfieldBase
-    
+
     def get_internal_type(self):
         return "CharField"
-    
+
     def formfield(self, **kwargs):
         from django.contrib.localflavor.us.forms import USSocialSecurityNumberField
         defaults = {'form_class': USSocialSecurityNumberField}
@@ -336,7 +349,8 @@ def formfield(self, **kwargs):
             ],
             [],
             {
-                'cipher':('cipher_type', {}),
+                'cipher': ('cipher_type', {}),
+                'block_type': ('block_type', {}),
             },
         ),
     ], ["^django_fields\.fields\..+?Field"])

src/django_fields/tests.py

@@ -51,7 +51,7 @@ class PickleObject(models.Model):
     data = PickleField()
 
 
-class EmailObject(models.Model): 
+class EmailObject(models.Model):
     max_email = 255
     email = EncryptedEmailField()
 
@@ -87,7 +87,7 @@ def test_encryption(self):
         encrypted_password = self._get_encrypted_password(obj.id)
         self.assertNotEqual(encrypted_password, password)
         self.assertTrue(encrypted_password.startswith('$AES$'))
-        
+
     def test_encryption_w_cipher(self):
         """
         Test that the database values are actually encrypted when using
@@ -104,6 +104,22 @@ def test_encryption_w_cipher(self):
         self.assertNotEqual(encrypted_password, password)
         self.assertTrue(encrypted_password.startswith('$AES$MODE_CBC$'))
 
+    def test_multiple_encryption_w_cipher(self):
+        """
+        Test that a single field can be reused without error.
+        """
+        password = 'this is a password!!'
+        obj = CipherEncObject(password=password)
+        obj.save()
+        obj = CipherEncObject.objects.get(id=obj.id)
+        self.assertEqual(password, obj.password)
+
+        password = 'another password!!'
+        obj = CipherEncObject(password=password)
+        obj.save()
+        obj = CipherEncObject.objects.get(id=obj.id)
+        self.assertEqual(password, obj.password)
+
     def test_max_field_length(self):
         password = 'a' * EncObject.max_password
         obj = EncObject(password = password)
@@ -122,7 +138,7 @@ def test_UTF8(self):
         obj.save()
         obj = EncObject.objects.get(id=obj.id)
         self.assertEqual(password, obj.password)
-    
+
     def test_consistent_encryption(self):
         """
         The same password should not encrypt the same way twice.
@@ -134,7 +150,7 @@ def test_consistent_encryption(self):
         for pwd_length in range(1,21):  # 1-20 inclusive
             enc_pwd_1, enc_pwd_2 = self._get_two_passwords(pwd_length)
             self.assertNotEqual(enc_pwd_1, enc_pwd_2)
-    
+
     def test_minimum_padding(self):
         """
         There should always be at least two chars of padding.
@@ -163,7 +179,7 @@ def _get_encrypted_password(self, id):
         passwords = map(lambda x: x[0], cursor.fetchall())
         self.assertEqual(len(passwords), 1)  # only one
         return passwords[0]
-    
+
     def _get_encrypted_password_cipher(self, id):
         cursor = connection.cursor()
         cursor.execute("select password from django_fields_cipherencobject where id = %s", [id,])
@@ -352,7 +368,7 @@ def test_UTF8(self):
         obj.save()
         obj = EmailObject.objects.get(id=obj.id)
         self.assertEqual(email, obj.email)
-    
+
     def test_consistent_encryption(self):
         """
         The same password should not encrypt the same way twice.
@@ -364,7 +380,7 @@ def test_consistent_encryption(self):
         for email_length in range(1,21):  # 1-20 inclusive
             enc_email_1, enc_email_2 = self._get_two_emails(email_length)
             self.assertNotEqual(enc_email_1, enc_email_2)
-    
+
     def test_minimum_padding(self):
         """
         There should always be at least two chars of padding.