doris-mcp-server/test/security/test_data_masking.py

#!/usr/bin/env python3
"""
Data masking tests
"""

import pytest

from doris_mcp_server.utils.security import (
    DataMaskingProcessor,
    AuthContext,
    SecurityLevel,
    MaskingRule
)


class TestDataMaskingProcessor:
    """Data masking processor tests"""

    @pytest.fixture
    def masking_processor(self, test_config):
        """Create data masking processor instance"""
        return DataMaskingProcessor(test_config)

    @pytest.fixture
    def internal_user_context(self):
        """Create internal user auth context"""
        return AuthContext(
            user_id="internal_user",
            roles=["data_analyst"],
            permissions=["read_data"],
            session_id="session_123",
            security_level=SecurityLevel.INTERNAL
        )

    @pytest.fixture
    def admin_context(self):
        """Create admin auth context"""
        return AuthContext(
            user_id="admin",
            roles=["data_admin"],
            permissions=["admin"],
            session_id="session_456",
            security_level=SecurityLevel.SECRET
        )

    @pytest.mark.asyncio
    async def test_phone_masking_for_internal_user(self, masking_processor, internal_user_context, sample_data):
        """Test phone number masking for internal user"""
        result = await masking_processor.process(sample_data, internal_user_context)
        
        # Phone numbers should be masked
        assert result[0]["phone"] == "138****5678"
        assert result[1]["phone"] == "139****4321"

    @pytest.mark.asyncio
    async def test_email_masking_for_internal_user(self, masking_processor, internal_user_context, sample_data):
        """Test email masking for internal user"""
        result = await masking_processor.process(sample_data, internal_user_context)
        
        # Emails should be masked
        assert result[0]["email"] == "z******n@example.com"
        assert result[1]["email"] == "l**i@example.com"

    @pytest.mark.asyncio
    async def test_no_masking_for_admin(self, masking_processor, admin_context, sample_data):
        """Test no masking for admin user"""
        result = await masking_processor.process(sample_data, admin_context)
        
        # Admin should see original data
        assert result[0]["phone"] == "13812345678"
        assert result[0]["email"] == "zhangsan@example.com"
        assert result[1]["phone"] == "13987654321"
        assert result[1]["email"] == "lisi@example.com"

    @pytest.mark.asyncio
    async def test_id_card_masking_for_confidential_data(self, masking_processor, internal_user_context, sample_data):
        """Test ID card masking for confidential data"""
        # Internal user should not see ID card details (confidential level)
        result = await masking_processor.process(sample_data, internal_user_context)
        
        # ID cards should be masked for internal users
        assert result[0]["id_card"] == "110101********1234"
        assert result[1]["id_card"] == "110101********2345"

    @pytest.mark.asyncio
    async def test_empty_data_handling(self, masking_processor, internal_user_context):
        """Test empty data handling"""
        empty_data = []
        
        result = await masking_processor.process(empty_data, internal_user_context)
        
        assert result == []

    @pytest.mark.asyncio
    async def test_null_value_handling(self, masking_processor, internal_user_context):
        """Test null value handling"""
        data_with_nulls = [
            {
                "id": 1,
                "name": "张三",
                "phone": None,
                "email": None,
                "id_card": None
            }
        ]
        
        result = await masking_processor.process(data_with_nulls, internal_user_context)
        
        # Null values should remain null
        assert result[0]["phone"] is None
        assert result[0]["email"] is None
        assert result[0]["id_card"] is None

    def test_phone_masking_algorithm(self, masking_processor):
        """Test phone masking algorithm"""
        params = {"mask_char": "*", "keep_prefix": 3, "keep_suffix": 4}
        
        result = masking_processor._mask_phone("13812345678", params)
        
        assert result == "138****5678"

    def test_email_masking_algorithm(self, masking_processor):
        """Test email masking algorithm"""
        params = {"mask_char": "*"}
        
        result = masking_processor._mask_email("zhangsan@example.com", params)
        
        assert result == "z******n@example.com"

    def test_id_card_masking_algorithm(self, masking_processor):
        """Test ID card masking algorithm"""
        params = {"mask_char": "*", "keep_prefix": 6, "keep_suffix": 4}
        
        result = masking_processor._mask_id_card("110101199001011234", params)
        
        assert result == "110101********1234"

    def test_name_masking_algorithm(self, masking_processor):
        """Test name masking algorithm"""
        params = {"mask_char": "*"}
        
        # Test 2-character name
        result = masking_processor._mask_name("张三", params)
        assert result == "张*"
        
        # Test 3-character name
        result = masking_processor._mask_name("李小明", params)
        assert result == "李*明"

    def test_partial_masking_algorithm(self, masking_processor):
        """Test partial masking algorithm"""
        params = {"mask_char": "*", "mask_ratio": 0.5}
        
        result = masking_processor._mask_partial("1234567890", params)
        
        # Should mask middle 50% of the string
        assert "*" in result
        assert len(result) == 10

    def test_should_apply_rule_logic(self, masking_processor, internal_user_context, admin_context):
        """Test masking rule application logic"""
        rule = MaskingRule(
            column_pattern=r".*phone.*",
            algorithm="phone_mask",
            parameters={"mask_char": "*", "keep_prefix": 3, "keep_suffix": 4},
            security_level=SecurityLevel.INTERNAL
        )
        
        # Internal user should have rule applied
        assert masking_processor._should_apply_rule(rule, internal_user_context) is True
        
        # Admin should not have rule applied
        assert masking_processor._should_apply_rule(rule, admin_context) is False

    def test_get_applicable_rules(self, masking_processor, internal_user_context):
        """Test getting applicable rules"""
        rules = masking_processor._get_applicable_rules(internal_user_context)
        
        # Should return some rules for internal user
        assert len(rules) > 0
        assert all(isinstance(rule, MaskingRule) for rule in rules)