Merge pull request #153 from alphagov/scan-files-on-upload

Scan files on upload

.travis.yml

@@ -8,7 +8,10 @@ matrix:
 # Need mongodb for testing
 services: mongodb
 # command to install dependencies
-install: "pip install -q -r requirements_for_tests.txt --use-mirrors"
+install: 
+  - pip install -q -r requirements_for_tests.txt --use-mirrors
+  - sudo apt-get install clamav -y
+  - sudo freshclam
 # command to run tests
 script:
   - ./run_tests.sh

backdrop/write/admin_ui.py

@@ -8,6 +8,7 @@
 from backdrop.core.upload.filters import first_sheet_filter
 from backdrop.write.signonotron2 import Signonotron2
 from backdrop.write.uploaded_file import UploadedFile, FileUploadException
+from backdrop.write.scanned_file import VirusSignatureError
 from ..core import cache_control
 
 
@@ -143,7 +144,10 @@ def _store_data(bucket_config):
             bucket = Bucket(db, bucket_config)
             upload.save(bucket, parser)
             return render_template('upload_ok.html')
-        except (FileUploadException, ParseError, ValidationError) as e:
+        except (VirusSignatureError,
+                FileUploadException,
+                ParseError,
+                ValidationError) as e:
             message = e.message
             app.logger.error(message)
             return _invalid_upload(message)

backdrop/write/scanned_file.py

@@ -0,0 +1,37 @@
+import hashlib
+import subprocess
+import os
+
+
+class VirusSignatureError(StandardError):
+    def __init__(self, message):
+        self.message = message
+
+
+class ScannedFile(object):
+    def __init__(self, file_object):
+        self.file_object = file_object
+        self._virus_signature = False
+        self._file_path = '/tmp/{0}'.format(
+            os.path.basename(self.file_object.filename))
+
+    @property
+    def has_virus_signature(self):
+        self._save_file_to_disk()
+        self._scan_file()
+        self._clean_up()
+        return self._virus_signature
+
+    def _save_file_to_disk(self):
+        self.file_object.save(self._file_path)
+
+    def _scan_file(self):
+        self._virus_signature = (self._virus_signature or
+                                 bool(subprocess.call(["clamscan",
+                                      self._file_path])))
+
+    def _clean_up(self):
+        # Remove temporary file
+        os.remove(self._file_path)
+        # Reset stream position on file_object so that it can be read again
+        self.file_object.seek(0)

backdrop/write/uploaded_file.py

@@ -1,3 +1,6 @@
+from backdrop.write.scanned_file import ScannedFile, VirusSignatureError
+
+
 class FileUploadException(IOError):
     def __init__(self, message):
         self.message = message
@@ -29,12 +32,20 @@ def _is_content_type_valid(self):
     def save(self, bucket, parser):
         if not self.valid:
             self.file_stream().close()
-            raise FileUploadException('Invalid file upload %s' %
-                                      self.file_object)
+            raise FileUploadException('Invalid file upload {0}'
+                                      .format(self.file_object.filename))
+        self.perform_virus_scan()
         data = parser(self.file_stream())
         bucket.parse_and_store(data)
         self.file_stream().close()
 
+    def perform_virus_scan(self):
+        if ScannedFile(self.file_object).has_virus_signature:
+            self.file_stream().close()
+            raise VirusSignatureError(
+                'File {0} could not be uploaded as it may contain a virus.'
+                .format(self.file_object.filename))
+
     @property
     def valid(self):
         return self._is_size_valid() and self._is_content_type_valid()

tests/write/test_scanned_file.py

@@ -0,0 +1,46 @@
+import unittest
+from hamcrest import assert_that, is_
+from mock import Mock
+from backdrop.write.scanned_file import ScannedFile
+import subprocess
+import os
+from tests.support.file_upload_test_case import FileUploadTestCase
+
+
+class TestScannedFile(FileUploadTestCase):
+
+    def setUp(self):
+        self.file_object = self._file_storage_wrapper("This is a test", "abc.txt")
+        self.scanner = ScannedFile(self.file_object)
+
+    def tearDown(self):
+        try:
+            os.remove('/tmp/abc.txt')
+        except OSError:
+            pass
+
+    def test_has_virus_signature(self):
+        self.scanner._virus_signature = True
+        self.scanner._save_file_to_disk = Mock()
+        self.scanner._scan_file = Mock()
+        self.scanner._clean_up = Mock()
+        assert_that(self.scanner.has_virus_signature, is_(True))
+        self.scanner._save_file_to_disk.assert_called_once_with()
+        self.scanner._scan_file.assert_called_once_with()
+        self.scanner._clean_up.assert_called_once_with()
+
+    def test_saving_a_file_to_disk(self):
+        self.scanner._save_file_to_disk()
+        assert_that(file('/tmp/abc.txt').read(), is_("This is a test"))
+
+    def test_cleaning_up_after_scanning(self):
+        self.file_object.save('/tmp/abc.txt')
+        self.scanner._clean_up()
+        assert_that(os.path.exists('/tmp/abc.txt'), is_(False))
+
+    def test_scanning_a_file(self):
+        mock_call = Mock(return_value = True)
+        subprocess.call = mock_call
+        self.scanner._scan_file()
+        assert_that(self.scanner._virus_signature, is_(True))
+        mock_call.assert_called_once_with(["clamscan", "/tmp/abc.txt"])

tests/write/test_uploaded_file.py

@@ -1,6 +1,7 @@
 from hamcrest import assert_that, is_
-from mock import Mock
+from mock import Mock, patch
 from backdrop.write.uploaded_file import UploadedFile, FileUploadException
+from backdrop.write.scanned_file import ScannedFile, VirusSignatureError
 from tests.support.file_upload_test_case import FileUploadTestCase
 
 
@@ -17,7 +18,7 @@ def test_files_under_1000000_octets_are_valid(self):
 
         assert_that(upload.valid, is_(True))
 
-    def test_files_under_1000000_octets_are_valid(self):
+    def test_files_over_1000000_octets_are_valid(self):
         upload = UploadedFile(self._file_storage_wrapper(
             "foo",
             content_type="text/csv",
@@ -32,6 +33,7 @@ def test_saving_file(self):
         ))
         bucket = Mock()
         parser = Mock()
+        upload.perform_virus_scan = Mock()
         parser.return_value = "some data"
         upload.save(bucket, parser)
         assert_that(parser.called, is_(True))
@@ -90,3 +92,10 @@ def test_files_with_no_content_type_are_invalid(self):
             self._file_storage_wrapper('foo', content_type=None))
 
         assert_that(upload.valid, is_(False))
+
+    @patch('backdrop.write.scanned_file.ScannedFile.has_virus_signature')
+    def test_perform_virus_scan(self, has_virus_signature):
+        file_storage_wrapper = self._file_storage_wrapper('foo', content_type=None)
+        upload = UploadedFile(file_storage_wrapper)
+        has_virus_signature.return_value = True
+        self.assertRaises(VirusSignatureError, upload.perform_virus_scan)