diff --git a/README.md b/README.md index 8906951..9ca4aff 100644 --- a/README.md +++ b/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 `_`; 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 diff --git a/docs/getting_started.rst b/docs/getting_started.rst index e4ad4d7..7e63a9a 100644 --- a/docs/getting_started.rst +++ b/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 ``_``; 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: diff --git a/succulent/processing.py b/succulent/processing.py index 840062e..c37bfe4 100644 --- a/succulent/processing.py +++ b/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) diff --git a/tests/test_succulent.py b/tests/test_succulent.py index a53bbb1..47f8dc7 100644 --- a/tests/test_succulent.py +++ b/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'')) + + 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'', 'x.html', 'File type ".html" is not allowed'), + 'script as jpg': (b'', '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.