fix: check type and content of uploaded images

README.md

@@ -202,8 +202,11 @@ To collect images, create a `configuration.yml` file in the root directory and d
 ```yml
 data:
   - key: # Key in POST request
+    allowed_extensions: ['jpg', 'png'] # Optional, see below
 ```
 
+Only images are accepted. By default, the allowed extensions are `jpg`, `jpeg`, `png`, `gif`, `bmp`, `webp`, `tif` and `tiff`; use `allowed_extensions` to narrow or extend this list. For these image types, the content of the file is checked as well, so a file renamed to `.jpg` that is not a JPEG image is rejected. Other extensions you allow (e.g. `heic`) are checked by extension only. A file without an extension (e.g. from an ESP32-CAM) is accepted if its content is an allowed image, and the matching extension is added. Rejected uploads return a `400` response with the reason. Images are stored as `<timestamp>_<filename>`; an image with the same name in the same second gets a counter (e.g. `_1`) instead of overwriting the earlier one.
+
 To store data collection timestamps, define the following setting in the `configuration.yml` file in the root directory:
 ```yml
 timestamp: true # false by default

docs/getting_started.rst

@@ -74,6 +74,9 @@ To collect images, create a ``configuration.yml`` file in the root directory and
 
     data:
       - key: # Key in POST request
+        allowed_extensions: ['jpg', 'png'] # Optional, see below
+
+Only images are accepted. By default, the allowed extensions are ``jpg``, ``jpeg``, ``png``, ``gif``, ``bmp``, ``webp``, ``tif`` and ``tiff``; use ``allowed_extensions`` to narrow or extend this list. For these image types, the content of the file is checked as well, so a file renamed to ``.jpg`` that is not a JPEG image is rejected. Other extensions you allow (e.g. ``heic``) are checked by extension only. A file without an extension (e.g. from an ESP32-CAM) is accepted if its content is an allowed image, and the matching extension is added. Rejected uploads return a ``400`` response with the reason. Images are stored as ``<timestamp>_<filename>``; an image with the same name in the same second gets a counter (e.g. ``_1``) instead of overwriting the earlier one.
 
 To store data collection timestamps, define the following setting in the ``configuration.yml`` file in the root directory:
 

succulent/processing.py

@@ -28,11 +28,27 @@ class Processing:
         enable_export (bool): Enable collected data export via the API.
         df (pandas.DataFrame): The DataFrame to hold the data (for 'csv', 'json', and 'sqlite' formats).
         key (str): The key to retrieve the image file from the request (for 'image' format).
+        allowed_extensions (list): File extensions accepted for uploaded images (for 'image' format).
 
     Raises:
         ValueError: Invalid format is provided.
     """
 
+    # Image types recognised by their content (file signature), mapped to
+    # the file extensions that may be used for them.
+    IMAGE_TYPES = {
+        'jpeg': ('jpg', 'jpeg'),
+        'png': ('png',),
+        'gif': ('gif',),
+        'bmp': ('bmp',),
+        'webp': ('webp',),
+        'tiff': ('tif', 'tiff'),
+    }
+
+    # Extensions accepted for the 'image' format when 'allowed_extensions'
+    # is not set in the configuration.
+    DEFAULT_IMAGE_EXTENSIONS = ['jpg', 'jpeg', 'png', 'gif', 'bmp', 'webp', 'tif', 'tiff']
+
     def __init__(self, format, config, unittest=False):
         # Validate format
         if format not in ['csv', 'json', 'sqlite', 'image', 'xml']:
@@ -91,6 +107,15 @@ def __init__(self, format, config, unittest=False):
         if format == 'image':
             self.key = config['data'][0]['key']
 
+            allowed = config['data'][0].get(
+                'allowed_extensions', self.DEFAULT_IMAGE_EXTENSIONS)
+            if (not isinstance(allowed, list) or not allowed
+                    or not all(isinstance(ext, str) and ext.strip('. ') for ext in allowed)):
+                raise ValueError(
+                    '"allowed_extensions" must be a non-empty list of file extensions, '
+                    'e.g. [\'jpg\', \'png\'].')
+            self.allowed_extensions = [ext.strip().lstrip('.').lower() for ext in allowed]
+
     @staticmethod
     def valid_xml_name(name):
         """Check whether a column name can be used as an XML element name.
@@ -105,6 +130,85 @@ def valid_xml_name(name):
                 and re.fullmatch(r'[A-Za-z_][\w.\-]*', name, flags=re.ASCII) is not None
                 and not name.lower().startswith('xml'))
 
+    @classmethod
+    def image_type(cls, header):
+        """Detect the image type from the first bytes of a file.
+
+        Args:
+            header (bytes): The first bytes (at least 12) of the file.
+
+        Returns:
+            str or None: One of the keys of ``IMAGE_TYPES``, or None if the
+            content is not a recognised image.
+        """
+        if header.startswith(b'\xff\xd8\xff'):
+            return 'jpeg'
+        if header.startswith(b'\x89PNG\r\n\x1a\n'):
+            return 'png'
+        if header[:6] in (b'GIF87a', b'GIF89a'):
+            return 'gif'
+        if header.startswith(b'BM'):
+            return 'bmp'
+        if header[:4] == b'RIFF' and header[8:12] == b'WEBP':
+            return 'webp'
+        if header[:4] in (b'II*\x00', b'MM\x00*'):
+            return 'tiff'
+        return None
+
+    def check_image(self, file, filename):
+        """Check that an uploaded file is an allowed image.
+
+        The extension must be in ``allowed_extensions``. For known image
+        types (see ``IMAGE_TYPES``), the content must also match: a script
+        renamed to ``.jpg`` is rejected. Extensions without a known
+        signature (e.g. ``heic``), if allowed in the configuration, are
+        checked by extension only. A file without an extension is accepted
+        if its content is an allowed image type, and the matching extension
+        is added (some cameras, e.g. ESP32-CAM, send no filename).
+
+        Args:
+            file (FileStorage): The uploaded file.
+            filename (str): The sanitised filename.
+
+        Raises:
+            ValueError: The file is empty, its extension is not allowed, or
+            its content does not match an allowed image type.
+
+        Returns:
+            str: The filename to store the file under.
+        """
+        header = file.stream.read(16)
+        file.stream.seek(0)
+        if not header:
+            raise ValueError('The uploaded file is empty.')
+
+        detected = self.image_type(header)
+        name, dot, extension = filename.rpartition('.')
+        extension = extension.lower() if dot and name else ''
+
+        allowed = ', '.join(self.allowed_extensions)
+
+        # No extension: name the file after its detected content
+        if not extension:
+            for candidate in self.IMAGE_TYPES.get(detected, ()):
+                if candidate in self.allowed_extensions:
+                    return f'{filename}.{candidate}'
+            raise ValueError(
+                f'The uploaded file is not an allowed image (allowed: {allowed}).')
+
+        if extension not in self.allowed_extensions:
+            raise ValueError(
+                f'File type ".{extension}" is not allowed (allowed: {allowed}).')
+
+        # Known image extension: the content must be of that type
+        expected = [kind for kind, extensions in self.IMAGE_TYPES.items()
+                    if extension in extensions]
+        if expected and detected not in expected:
+            raise ValueError(
+                f'The content of the uploaded file does not match its extension (.{extension}).')
+
+        return filename
+
     def parameters(self):
         """Generate URL parameters based on data columns.
 
@@ -364,12 +468,17 @@ def process(self, req):
 
         if self.format == 'image':
             # Retrieve image from request
-            file = req.files[self.key]
+            file = req.files.get(self.key)
+            if file is None:
+                raise ValueError(f'No file was uploaded in the "{self.key}" field.')
 
             filename = secure_filename(file.filename or '')
             if not filename:
                 filename = 'upload'
 
+            # Allowed extension and matching content
+            filename = self.check_image(file, filename)
+
             # Timestamp
             timestamp = now.strftime('%Y-%m-%d_%H%M%S')
 
@@ -382,6 +491,14 @@ def process(self, req):
             if os.path.commonpath([directory, destination]) != directory:
                 raise ValueError('Invalid file path.')
 
+            # Do not overwrite an image uploaded with the same name in the
+            # same second; add a counter instead (e.g. ..._photo_1.jpg)
+            base, extension = os.path.splitext(destination)
+            counter = 1
+            while os.path.exists(destination):
+                destination = f'{base}_{counter}{extension}'
+                counter += 1
+
             # Store image to device
             file.save(destination)
 

tests/test_succulent.py

@@ -1,6 +1,7 @@
 from succulent import __version__
 import pandas as pd
 import unittest
+import io
 import json
 import os
 import shutil
@@ -150,8 +151,9 @@ def setUp(self):
         self.processing = Processing(config=config, format='image', unittest=True)
 
     def test_process_image(self):
-        # Temporary file to simulate an image
-        image_data = b'Test image data'
+        # Temporary file to simulate an image (starts with the JPEG signature,
+        # since the content of uploaded images is checked)
+        image_data = b'\xff\xd8\xff\xe0' + b'Test image data'
         with open('temp_image.jpg', 'wb') as f:
             f.write(image_data)
 
@@ -570,6 +572,105 @@ def test_non_object_json_returns_400(self):
                                  'Invalid JSON: the request body must be a JSON object.')
 
 
+class TestImageUpload(unittest.TestCase):
+    """
+    Test case for checking uploaded images.
+    """
+
+    JPEG = b'\xff\xd8\xff\xe0' + b'\x00' * 32
+    PNG = b'\x89PNG\r\n\x1a\n' + b'\x00' * 32
+
+    def setUp(self):
+        self.config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'image.yml')
+        self.api = SucculentAPI(host='0.0.0.0', port=8080, config=self.config_path, format='image')
+        self.client = self.api.app.test_client()
+
+    def tearDown(self):
+        shutil.rmtree(self.api.processing.directory, ignore_errors=True)
+
+    def upload(self, content, filename, field='image', client=None):
+        client = client or self.client
+        return client.post('/measure', data={field: (io.BytesIO(content), filename)},
+                           content_type='multipart/form-data')
+
+    def stored_files(self):
+        return sorted(os.listdir(self.api.processing.directory))
+
+    def test_signatures(self):
+        self.assertEqual(Processing.image_type(self.JPEG), 'jpeg')
+        self.assertEqual(Processing.image_type(self.PNG), 'png')
+        self.assertEqual(Processing.image_type(b'GIF89a' + b'\x00' * 10), 'gif')
+        self.assertEqual(Processing.image_type(b'BM' + b'\x00' * 10), 'bmp')
+        self.assertEqual(Processing.image_type(b'RIFF\x00\x00\x00\x00WEBP'), 'webp')
+        self.assertEqual(Processing.image_type(b'II*\x00' + b'\x00' * 8), 'tiff')
+        self.assertEqual(Processing.image_type(b'MM\x00*' + b'\x00' * 8), 'tiff')
+        self.assertIsNone(Processing.image_type(b'<?php echo 1; ?>'))
+
+    def test_valid_images_are_stored(self):
+        for content, filename in [(self.JPEG, 'a.jpg'), (self.JPEG, 'b.JPEG'), (self.PNG, 'c.png')]:
+            with self.subTest(filename=filename):
+                self.assertEqual(self.upload(content, filename).status_code, 200)
+        self.assertEqual(len(self.stored_files()), 3)
+
+    def test_file_without_extension_gets_one(self):
+        # e.g. ESP32-CAM sends the image without a filename
+        self.assertEqual(self.upload(self.JPEG, '').status_code, 200)
+        self.assertEqual(self.upload(self.PNG, 'capture').status_code, 200)
+        files = self.stored_files()
+        self.assertTrue(any(f.endswith('_upload.jpg') for f in files), files)
+        self.assertTrue(any(f.endswith('_capture.png') for f in files), files)
+
+    def test_rejected_files(self):
+        cases = {
+            'executable': (b'MZ\x90\x00', 'x.exe', 'File type ".exe" is not allowed'),
+            'html': (b'<script>alert(1)</script>', 'x.html', 'File type ".html" is not allowed'),
+            'script as jpg': (b'<?php system($_GET[1]); ?>', 'x.jpg', 'does not match its extension'),
+            'png as jpg': (self.PNG, 'x.jpg', 'does not match its extension'),
+            'unknown, no extension': (b'hello', 'capture', 'not an allowed image'),
+            'empty': (b'', 'x.png', 'The uploaded file is empty.'),
+        }
+        for name, (content, filename, message) in cases.items():
+            with self.subTest(case=name):
+                response = self.upload(content, filename)
+                self.assertEqual(response.status_code, 400)
+                self.assertIn(message, response.json['message'])
+        self.assertFalse(os.path.exists(self.api.processing.directory)
+                         and self.stored_files())
+
+    def test_missing_file_field(self):
+        for data in [{}, {'photo': (io.BytesIO(self.PNG), 'a.png')}]:
+            with self.subTest(data=list(data)):
+                response = self.client.post('/measure', data=data,
+                                            content_type='multipart/form-data')
+                self.assertEqual(response.status_code, 400)
+                self.assertEqual(response.json['message'],
+                                 'No file was uploaded in the "image" field.')
+
+    def test_same_name_is_not_overwritten(self):
+        for _ in range(3):
+            self.assertEqual(self.upload(self.JPEG, 'photo.jpg').status_code, 200)
+        self.assertEqual(len(self.stored_files()), 3)
+
+    def test_allowed_extensions_from_config(self):
+        config = {'data': [{'key': 'image', 'allowed_extensions': ['.PNG', 'heic']}]}
+        processing = Processing(config=config, format='image', unittest=True)
+        self.assertEqual(processing.allowed_extensions, ['png', 'heic'])
+
+        self.api.processing.allowed_extensions = processing.allowed_extensions
+        self.assertEqual(self.upload(self.PNG, 'a.png').status_code, 200)
+        # Not in the list anymore
+        self.assertEqual(self.upload(self.JPEG, 'a.jpg').status_code, 400)
+        # No known signature: checked by extension only
+        self.assertEqual(self.upload(b'\x00\x00\x00\x18ftypheic', 'a.heic').status_code, 200)
+
+    def test_invalid_allowed_extensions(self):
+        for value in [[], 'jpg', [''], [1]]:
+            with self.subTest(value=value):
+                with self.assertRaises(ValueError):
+                    Processing(config={'data': [{'key': 'image', 'allowed_extensions': value}]},
+                               format='image', unittest=True)
+
+
 class TestStaleData(unittest.TestCase):
     """
     Test that /data and /export do not return stale in-memory data.