mcp/tests/test_security.py

import unittest

from utils.security import (
    generate_gateway_session_id,
    hash_gateway_session_id,
    constant_time_equals
)


class TestSecurity(unittest.TestCase):
    def test_generate_gateway_session_id_format(self):
        session_id = generate_gateway_session_id()

        # Should start with GWS_
        self.assertTrue(session_id.startswith('GWS_'))

        # Should be long enough (GWS_ + 32+ chars)
        self.assertGreater(len(session_id), 36)

    def test_generate_gateway_session_id_uniqueness(self):
        # Generate multiple IDs and ensure they're all different
        ids = [generate_gateway_session_id() for _ in range(100)]
        unique_ids = set(ids)

        self.assertEqual(len(ids), len(unique_ids))

    def test_hash_gateway_session_id_consistent(self):
        session_id = 'GWS_test123'

        hash1 = hash_gateway_session_id(session_id)
        hash2 = hash_gateway_session_id(session_id)

        self.assertEqual(hash1, hash2)

    def test_hash_gateway_session_id_different_for_different_inputs(self):
        session_id_1 = 'GWS_test123'
        session_id_2 = 'GWS_test456'

        hash1 = hash_gateway_session_id(session_id_1)
        hash2 = hash_gateway_session_id(session_id_2)

        self.assertNotEqual(hash1, hash2)

    def test_hash_gateway_session_id_length(self):
        session_id = 'GWS_xyz'
        hashed = hash_gateway_session_id(session_id)

        # SHA256 produces 64 hex characters
        self.assertEqual(len(hashed), 64)

    def test_hash_gateway_session_id_raises_on_empty(self):
        with self.assertRaises(ValueError) as context:
            hash_gateway_session_id('')

        self.assertIn('gateway_session_id is required', str(context.exception))

    def test_hash_gateway_session_id_raises_on_none(self):
        with self.assertRaises(ValueError) as context:
            hash_gateway_session_id(None)

        self.assertIn('gateway_session_id is required', str(context.exception))

    def test_hash_gateway_session_id_handles_whitespace(self):
        # Hash should handle whitespace consistently
        session_id_with_spaces = '  GWS_test  '

        hash_result = hash_gateway_session_id(session_id_with_spaces)

        # Should be a valid hash
        self.assertEqual(len(hash_result), 64)

        # Same input should give same hash
        hash_result_2 = hash_gateway_session_id(session_id_with_spaces)
        self.assertEqual(hash_result, hash_result_2)

    def test_constant_time_equals_true_for_same_strings(self):
        result = constant_time_equals('test123', 'test123')
        self.assertTrue(result)

    def test_constant_time_equals_false_for_different_strings(self):
        result = constant_time_equals('test123', 'test456')
        self.assertFalse(result)

    def test_constant_time_equals_handles_none(self):
        result1 = constant_time_equals(None, None)
        self.assertTrue(result1)

        result2 = constant_time_equals('test', None)
        self.assertFalse(result2)

        result3 = constant_time_equals(None, 'test')
        self.assertFalse(result3)

    def test_constant_time_equals_handles_empty_strings(self):
        result = constant_time_equals('', '')
        self.assertTrue(result)

    def test_constant_time_equals_case_sensitive(self):
        result = constant_time_equals('Test', 'test')
        self.assertFalse(result)

    def test_generate_session_id_url_safe(self):
        # URL-safe base64 should not contain +, /, or =
        session_id = generate_gateway_session_id()
        token_part = session_id[4:]  # Remove 'GWS_' prefix

        self.assertNotIn('+', token_part)
        self.assertNotIn('/', token_part)
        # May or may not have = padding, but should be URL-safe

    def test_hash_deterministic(self):
        # Same input should always produce same output
        session_id = 'GWS_deterministic_test'
        hashes = [hash_gateway_session_id(session_id) for _ in range(10)]

        self.assertEqual(len(set(hashes)), 1)


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