AWS Model Context Protocol Server

"""Unit tests for the security module.""" from unittest.mock import mock_open, patch import pytest import yaml from aws_mcp_server.security import ( DEFAULT_DANGEROUS_COMMANDS, DEFAULT_SAFE_PATTERNS, SecurityConfig, ValidationRule, check_regex_rules, is_service_command_safe, load_security_config, reload_security_config, validate_aws_command, validate_command, validate_pipe_command, ) def test_is_service_command_safe(): """Test the is_service_command_safe function.""" # Test with known safe pattern assert is_service_command_safe("aws s3 ls", "s3") is True # Test with known dangerous pattern that has safe override assert is_service_command_safe("aws s3 ls --profile test", "s3") is True # Test with known dangerous pattern with no safe override assert is_service_command_safe("aws s3 rb s3://my-bucket", "s3") is False # Test with unknown service assert is_service_command_safe("aws unknown-service command", "unknown-service") is False def test_check_regex_rules(): """Test the check_regex_rules function.""" # Test with a pattern that should match with patch("aws_mcp_server.security.SECURITY_CONFIG") as mock_config: mock_config.regex_rules = { "general": [ ValidationRule( pattern=r"aws .* --profile\s+(root|admin|administrator)", description="Prevent use of sensitive profiles", error_message="Using sensitive profiles (root, admin) is restricted", regex=True, ) ] } # Should match the rule error = check_regex_rules("aws s3 ls --profile root") assert error is not None assert "Using sensitive profiles" in error # Should not match assert check_regex_rules("aws s3 ls --profile user") is None @patch("aws_mcp_server.security.SECURITY_MODE", "strict") def test_validate_aws_command_basic(): """Test basic validation of AWS commands.""" # Valid command should not raise validate_aws_command("aws s3 ls") # Invalid commands should raise ValueError with pytest.raises(ValueError, match="Commands must start with 'aws'"): validate_aws_command("s3 ls") with pytest.raises(ValueError, match="must include an AWS service"): validate_aws_command("aws") @patch("aws_mcp_server.security.SECURITY_MODE", "strict") def test_validate_aws_command_dangerous(): """Test validation of dangerous AWS commands.""" # Use a test config with patch("aws_mcp_server.security.SECURITY_CONFIG") as mock_config: mock_config.dangerous_commands = { "iam": ["aws iam create-user", "aws iam create-access-key"], "ec2": ["aws ec2 terminate-instances"], } mock_config.safe_patterns = { "iam": ["aws iam create-user --help"], "ec2": [], } mock_config.regex_rules = {} # Dangerous command should raise ValueError with pytest.raises(ValueError, match="restricted for security reasons"): validate_aws_command("aws iam create-user --user-name test-user") # Help on dangerous command should be allowed validate_aws_command("aws iam create-user --help") # Dangerous command with no safe override should raise with pytest.raises(ValueError, match="restricted for security reasons"): validate_aws_command("aws ec2 terminate-instances --instance-id i-12345") @patch("aws_mcp_server.security.SECURITY_MODE", "strict") def test_validate_aws_command_regex(): """Test validation of AWS commands with regex rules.""" # Set up command for testing profile_command = "aws s3 ls --profile root" policy_command = """aws s3api put-bucket-policy --bucket my-bucket --policy "{\\"Version\\":\\"2012-10-17\\",\ \\"Statement\\":[{\\"Effect\\":\\"Allow\\",\\"Principal\\":\\"*\\",\\"Action\\":\\"s3:GetObject\\",\ \\"Resource\\":\\"arn:aws:s3:::my-bucket/*\\"}]}" """ # We need to patch both the check_regex_rules function and the config with patch("aws_mcp_server.security.SECURITY_CONFIG") as mock_config: mock_config.dangerous_commands = {} mock_config.safe_patterns = {} # Test for the root profile check with patch("aws_mcp_server.security.check_regex_rules") as mock_check: mock_check.return_value = "Using sensitive profiles is restricted" with pytest.raises(ValueError, match="Using sensitive profiles is restricted"): validate_aws_command(profile_command) # Verify check_regex_rules was called mock_check.assert_called_once() # Test for the bucket policy check with patch("aws_mcp_server.security.check_regex_rules") as mock_check: # Have the mock return error for the policy command mock_check.return_value = "Creating public bucket policies is restricted" with pytest.raises(ValueError, match="Creating public bucket policies is restricted"): validate_aws_command(policy_command) # Verify check_regex_rules was called mock_check.assert_called_once() @patch("aws_mcp_server.security.SECURITY_MODE", "permissive") def test_validate_aws_command_permissive(): """Test validation of AWS commands in permissive mode.""" # In permissive mode, dangerous commands should be allowed with patch("aws_mcp_server.security.logger.warning") as mock_warning: validate_aws_command("aws iam create-user --user-name test-user") mock_warning.assert_called_once() @patch("aws_mcp_server.security.SECURITY_MODE", "strict") def test_validate_pipe_command(): """Test validation of piped commands.""" # Mock the validate_aws_command and validate_unix_command functions with patch("aws_mcp_server.security.validate_aws_command") as mock_aws_validate: with patch("aws_mcp_server.security.validate_unix_command") as mock_unix_validate: # Set up return values mock_unix_validate.return_value = True # Test valid piped command validate_pipe_command("aws s3 ls | grep bucket") mock_aws_validate.assert_called_once_with("aws s3 ls") # Reset mocks mock_aws_validate.reset_mock() mock_unix_validate.reset_mock() # Test command with unrecognized Unix command mock_unix_validate.return_value = False with pytest.raises(ValueError, match="not allowed"): validate_pipe_command("aws s3 ls | unknown_command") # Empty command should raise with pytest.raises(ValueError, match="Empty command"): validate_pipe_command("") # Empty second command test # Configure split_pipe_command to return a list with an empty second command with patch("aws_mcp_server.security.split_pipe_command") as mock_split_pipe: mock_split_pipe.return_value = ["aws s3 ls", ""] with pytest.raises(ValueError, match="Empty command at position"): validate_pipe_command("aws s3 ls | ") @patch("aws_mcp_server.security.SECURITY_MODE", "strict") def test_validate_command(): """Test the centralized validate_command function.""" # Simple AWS command validate_command("aws s3 ls") # Piped command validate_command("aws s3 ls | grep bucket") # Invalid command with pytest.raises(ValueError): validate_command("s3 ls") def test_load_security_config_default(): """Test loading security configuration with defaults.""" with patch("aws_mcp_server.security.SECURITY_CONFIG_PATH", ""): config = load_security_config() # Should have loaded default values assert config.dangerous_commands == DEFAULT_DANGEROUS_COMMANDS assert config.safe_patterns == DEFAULT_SAFE_PATTERNS # Should have regex rules converted from DEFAULT_REGEX_RULES assert "general" in config.regex_rules assert len(config.regex_rules["general"]) > 0 assert isinstance(config.regex_rules["general"][0], ValidationRule) def test_load_security_config_custom(): """Test loading security configuration from a custom file.""" # Mock YAML file contents test_config = { "dangerous_commands": {"test_service": ["aws test_service dangerous_command"]}, "safe_patterns": {"test_service": ["aws test_service safe_pattern"]}, "regex_rules": {"test_service": [{"pattern": "test_pattern", "description": "Test description", "error_message": "Test error message"}]}, } # Mock the open function to return our test config with patch("builtins.open", mock_open(read_data=yaml.dump(test_config))): with patch("aws_mcp_server.security.SECURITY_CONFIG_PATH", "/fake/path.yaml"): with patch("pathlib.Path.exists", return_value=True): config = load_security_config() # Should have our custom values assert "test_service" in config.dangerous_commands assert "test_service" in config.safe_patterns assert "test_service" in config.regex_rules assert config.regex_rules["test_service"][0].pattern == "test_pattern" def test_load_security_config_error(): """Test error handling when loading security configuration.""" with patch("builtins.open", side_effect=Exception("Test error")): with patch("aws_mcp_server.security.SECURITY_CONFIG_PATH", "/fake/path.yaml"): with patch("pathlib.Path.exists", return_value=True): with patch("aws_mcp_server.security.logger.error") as mock_error: with patch("aws_mcp_server.security.logger.warning") as mock_warning: config = load_security_config() # Should log error and warning mock_error.assert_called_once() mock_warning.assert_called_once() # Should still have default values assert config.dangerous_commands == DEFAULT_DANGEROUS_COMMANDS def test_reload_security_config(): """Test reloading security configuration.""" with patch("aws_mcp_server.security.load_security_config") as mock_load: mock_load.return_value = SecurityConfig(dangerous_commands={"test": ["test"]}, safe_patterns={"test": ["test"]}) reload_security_config() # Should have called load_security_config mock_load.assert_called_once() # Integration-like tests for specific dangerous commands @patch("aws_mcp_server.security.SECURITY_MODE", "strict") def test_specific_dangerous_commands(): """Test validation of specific dangerous commands.""" # Configure the SECURITY_CONFIG with some dangerous commands with patch("aws_mcp_server.security.SECURITY_CONFIG") as mock_config: mock_config.dangerous_commands = { "iam": ["aws iam create-user", "aws iam create-access-key", "aws iam attach-user-policy"], "ec2": ["aws ec2 terminate-instances"], "s3": ["aws s3 rb"], "rds": ["aws rds delete-db-instance"], } mock_config.safe_patterns = { "iam": ["aws iam get-", "aws iam list-"], "ec2": ["aws ec2 describe-"], "s3": ["aws s3 ls"], "rds": ["aws rds describe-"], } mock_config.regex_rules = {} # IAM dangerous commands with pytest.raises(ValueError, match="restricted for security reasons"): validate_aws_command("aws iam create-user --user-name test-user") with pytest.raises(ValueError, match="restricted for security reasons"): validate_aws_command("aws iam create-access-key --user-name test-user") with pytest.raises(ValueError, match="restricted for security reasons"): validate_aws_command("aws iam attach-user-policy --user-name test-user --policy-arn arn:aws:iam::aws:policy/AdministratorAccess") # EC2 dangerous commands with pytest.raises(ValueError, match="restricted for security reasons"): validate_aws_command("aws ec2 terminate-instances --instance-ids i-12345") # S3 dangerous commands with pytest.raises(ValueError, match="restricted for security reasons"): validate_aws_command("aws s3 rb s3://my-bucket --force") # RDS dangerous commands with pytest.raises(ValueError, match="restricted for security reasons"): validate_aws_command("aws rds delete-db-instance --db-instance-identifier my-db --skip-final-snapshot") # Tests for safe patterns overriding dangerous commands @patch("aws_mcp_server.security.SECURITY_MODE", "strict") def test_safe_overrides(): """Test safe patterns that override dangerous commands.""" # IAM help commands should be allowed even if potentially dangerous validate_aws_command("aws iam --help") validate_aws_command("aws iam help") validate_aws_command("aws iam get-user --user-name test-user") validate_aws_command("aws iam list-users") # EC2 describe commands should be allowed validate_aws_command("aws ec2 describe-instances") # S3 list commands should be allowed validate_aws_command("aws s3 ls") validate_aws_command("aws s3api list-buckets") # Tests for complex regex patterns @patch("aws_mcp_server.security.SECURITY_MODE", "strict") def test_complex_regex_patterns(): """Test more complex regex patterns.""" # Instead of testing the regex directly, test the behavior we expect dangerous_sg_command = "aws ec2 authorize-security-group-ingress --group-id sg-12345 --protocol tcp --port 22 --cidr 0.0.0.0/0" safe_sg_command_80 = "aws ec2 authorize-security-group-ingress --group-id sg-12345 --protocol tcp --port 80 --cidr 0.0.0.0/0" # Define the validation rule directly ValidationRule( pattern=r"aws ec2 authorize-security-group-ingress.*--cidr\s+0\.0\.0\.0/0.*--port\s+(?!80|443)\d+", description="Prevent open security groups for non-web ports", error_message="Security group error", regex=True, ) # Test with mocked check_regex_rules with patch("aws_mcp_server.security.SECURITY_CONFIG") as mock_config: mock_config.dangerous_commands = {} mock_config.safe_patterns = {} with patch("aws_mcp_server.security.check_regex_rules") as mock_check: # Set up mock to return error for the dangerous command mock_check.side_effect = lambda cmd, svc=None: "Security group error" if "--port 22" in cmd else None # Test dangerous command raises error with pytest.raises(ValueError, match="Security group error"): validate_aws_command(dangerous_sg_command) # Test safe command doesn't raise mock_check.reset_mock() mock_check.return_value = None # Explicit safe return validate_aws_command(safe_sg_command_80) # Should not raise