test_security.py

import unittest
import app as app_module
from app import app
import io
import time

class SecurityTestCase(unittest.TestCase):
    def setUp(self):
        app.config['TESTING'] = True
        self.client = app.test_client()
        # Reset rate limiting counts for each test
        app_module.upload_counts.clear()
        # Mock model_loaded to True to bypass model initialization in tests
        app_module.model_loaded = True

    def test_download_valid_filename(self):
        # Even if the file doesn't exist, it should pass the regex and reach send_from_directory (404)
        valid_uuid = "a" * 32
        response = self.client.get(f'/download/colorized_{valid_uuid}.png')
        self.assertEqual(response.status_code, 404)

    def test_download_invalid_format(self):
        # Too short UUID
        response = self.client.get('/download/colorized_abc123.png')
        self.assertEqual(response.status_code, 400)
        self.assertIn(b'Invalid filename format', response.data)

        # Missing prefix
        valid_uuid = "a" * 32
        response = self.client.get(f'/download/{valid_uuid}.png')
        self.assertEqual(response.status_code, 400)

        # Invalid extension
        response = self.client.get(f'/download/colorized_{valid_uuid}.exe')
        self.assertEqual(response.status_code, 400)

    def test_path_traversal_prevention(self):
        # Path traversal with encoded slashes might still be caught by Flask or return 400 due to regex
        response = self.client.get('/download/..%2f..%2fapp.py')
        # If it reaches the handler, it should be 400. If Flask blocks it early, it might be 404.
        self.assertIn(response.status_code, [400, 404])

        # Test a case that DEFINITELY reaches the handler but is malicious
        response = self.client.get('/download/colorized_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.png%00')
        self.assertEqual(response.status_code, 400)

    def test_rate_limiting(self):
        # The limit is 10 uploads per 60 seconds
        for _ in range(10):
            data = {'file': (io.BytesIO(b"fake image data"), 'test.png')}
            response = self.client.post('/upload', data=data, content_type='multipart/form-data')
            # It might return 400 or 500 depending on model state, but shouldn't be 429 yet
            self.assertNotEqual(response.status_code, 429)

        # 11th request should be rate limited
        data = {'file': (io.BytesIO(b"fake image data"), 'test.png')}
        response = self.client.post('/upload', data=data, content_type='multipart/form-data')
        self.assertEqual(response.status_code, 429)
        self.assertIn(b'Rate limit exceeded', response.data)

    def test_upload_invalid_magic_bytes(self):
        # Test JPEG extension with plain text content
        data = {'file': (io.BytesIO(b"this is not a jpeg"), 'test.jpg')}
        response = self.client.post('/upload', data=data, content_type='multipart/form-data')
        self.assertEqual(response.status_code, 400)
        self.assertIn(b'Invalid image content', response.data)

        # Test PNG extension with plain text content
        data = {'file': (io.BytesIO(b"this is not a png"), 'test.png')}
        response = self.client.post('/upload', data=data, content_type='multipart/form-data')
        self.assertEqual(response.status_code, 400)
        self.assertIn(b'Invalid image content', response.data)

if __name__ == '__main__':
    unittest.main()