Check magic numbers for filetype

Author: hirunatan

testing/media/testfile/the_file.png

@@ -3,24 +3,32 @@
 from django.template.defaultfilters import filesizeformat
 from django.utils.translation import ugettext as _
 
+import magic
+
 class ValidatedFileField(models.FileField):
     def __init__(self, *args, **kwargs):
-        self.content_types = kwargs.pop("content_types")
-        self.max_upload_size = kwargs.pop("max_upload_size")
+        self.content_types = kwargs.pop("content_types", [])
+        self.max_upload_size = kwargs.pop("max_upload_size", 0)
         super(ValidatedFileField, self).__init__(*args, **kwargs)
 
     def clean(self, *args, **kwargs):        
         data = super(ValidatedFileField, self).clean(*args, **kwargs)
 
-        file = data.file
-        try:
-            content_type = file.content_type
-            if content_type in self.content_types:
-                if file._size > self.max_upload_size:
-                    raise forms.ValidationError(_('Please keep filesize under %s. Current filesize %s') % (filesizeformat(self.max_upload_size), filesizeformat(file._size)))
-            else:
+        if self.content_types:
+            file = data.file
+            content_type_headers = getattr(file, 'content_type', '')
+
+            mg = magic.Magic(mime = True)
+            content_type_magic = mg.from_buffer(file.read(1024))
+            file.seek(0)
+
+            if not content_type_headers in self.content_types or not content_type_magic in self.content_types:
                 raise forms.ValidationError(_('Filetype %s not supported.') % (file.content_type))
-        except AttributeError:
-            pass        
+
+        if self.max_upload_size:
+            if file._size > self.max_upload_size:
+                raise forms.ValidationError(_('Please keep filesize under %s. Current filesize %s') %
+                                            (filesizeformat(self.max_upload_size), filesizeformat(file._size)))
 
         return data
+

validatedfile/models.py

@@ -5,28 +5,14 @@
 
 import os.path
 
-from models import TestModel
-from forms import TestModelForm
+from models import TestModel, TestModelNoValidate
+from forms import TestModelForm, TestModelNoValidateForm
 
 class ValidatedFileFieldTest(TestCase):
 
     SAMPLE_FILES_PATH = 'validatedfile/tests/sample_files'
 
 
-    def _get_sample_file(self, filename):
-        path = os.path.join(self.SAMPLE_FILES_PATH, filename)
-        return open(path)
-
-
-    def _check_file_url(self, filefield, filename):
-        url = os.path.join(settings.MEDIA_URL, filefield.field.upload_to, filename)
-        self.assertEqual(filefield.url, url)
-
-        
-    def _get_file_url(self, filename):
-        return os.path.join(MEDIA_ROOT, prefix, filename)
-
-
     def test_create_empty_instance(self):
         instance = TestModel.objects.create()
         instance.save()
@@ -40,11 +26,12 @@ def test_create_instance_with_file(self):
 
         self._check_file_url(instance.the_file, 'the_file.png')
 
+        from ipdb import set_trace; set_trace()
         instance.the_file.delete()
         instance.delete()
 
 
-    def test_form(self):
+    def test_form_ok(self):
         uploaded_file = SimpleUploadedFile(
                 name = 'the_file.png',
                 content = self._get_sample_file('image2k.png').read(),
@@ -60,26 +47,80 @@ def test_form(self):
         instance.delete()
 
 
+    def test_form_invalid_size(self):
+        uploaded_file = SimpleUploadedFile(
+                name = 'the_file.png',
+                content = self._get_sample_file('image15k.png').read(),
+                content_type = 'image/png',
+            )
+        form = TestModelForm(data = {}, files = {'the_file': uploaded_file})
+        self.assertFalse(form.is_valid())
+        self.assertEqual(len(form.errors), 1)
+        self.assertEqual(len(form.errors['the_file']), 1)
+
+
     def test_form_invalid_filetype(self):
         uploaded_file = SimpleUploadedFile(
                 name = 'the_file.pdf',
-                content = self._get_sample_file('document.pdf').read(),
-                content_type = 'apllication/pdf',
+                content = self._get_sample_file('document1k.pdf').read(),
+                content_type = 'application/pdf',
             )
         form = TestModelForm(data = {}, files = {'the_file': uploaded_file})
         self.assertFalse(form.is_valid())
         self.assertEqual(len(form.errors), 1)
         self.assertEqual(len(form.errors['the_file']), 1)
 
 
-    def test_form_invalid_size(self):
+    def test_form_invalid_filetype_and_size(self):
         uploaded_file = SimpleUploadedFile(
                 name = 'the_file.pdf',
-                content = self._get_sample_file('image15k.png').read(),
+                content = self._get_sample_file('document15k.pdf').read(),
+                content_type = 'application/pdf',
+            )
+        form = TestModelForm(data = {}, files = {'the_file': uploaded_file})
+        self.assertFalse(form.is_valid())
+        self.assertEqual(len(form.errors), 1)
+        self.assertEqual(len(form.errors['the_file']), 1)
+
+
+    def test_form_fake_filetype(self):
+        uploaded_file = SimpleUploadedFile(
+                name = 'the_file.png',
+                content = self._get_sample_file('document1k.pdf').read(),
                 content_type = 'image/png',
             )
         form = TestModelForm(data = {}, files = {'the_file': uploaded_file})
         self.assertFalse(form.is_valid())
         self.assertEqual(len(form.errors), 1)
         self.assertEqual(len(form.errors['the_file']), 1)
 
+
+    def test_form_no_validate(self):
+        uploaded_file = SimpleUploadedFile(
+                name = 'the_file.pdf',
+                content = self._get_sample_file('document15k.pdf').read(),
+                content_type = 'application/pdf',
+            )
+        form = TestModelNoValidateForm(data = {}, files = {'the_file': uploaded_file})
+        self.assertTrue(form.is_valid())
+        instance = form.save()
+
+        self._check_file_url(instance.the_file, 'the_file.pdf')
+
+        instance.the_file.delete()
+        instance.delete()
+
+
+    def _get_sample_file(self, filename):
+        path = os.path.join(self.SAMPLE_FILES_PATH, filename)
+        return open(path)
+
+
+    def _check_file_url(self, filefield, filename):
+        url = os.path.join(settings.MEDIA_URL, filefield.field.upload_to, filename)
+        self.assertEqual(filefield.url, url)
+
+        
+    def _get_file_url(self, filename):
+        return os.path.join(MEDIA_ROOT, prefix, filename)
+

validatedfile/tests/__init__.py

@@ -1,8 +1,13 @@
 from django import forms
-from models import TestModel
+from models import TestModel, TestModelNoValidate
 
 class TestModelForm(forms.ModelForm):
 
     class Meta:
         model = TestModel
 
+class TestModelNoValidateForm(forms.ModelForm):
+    
+    class Meta:
+        model = TestModelNoValidate
+

validatedfile/tests/forms.py

@@ -10,3 +10,9 @@ class TestModel(models.Model):
                     content_types = ['image/png'],
                     max_upload_size = 10240)
 
+class TestModelNoValidate(models.Model):
+    the_file = ValidatedFileField(
+                    null = True,
+                    blank = True,
+                    upload_to = 'testfile')
+