atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_token_storage.py

@@ -0,0 +1,505 @@
+"""Unit tests for MCP token storage.
+
+Tests the secure per-user token storage module including:
+- Token encryption/decryption
+- Token storage and retrieval
+- Token expiration handling
+- Per-user isolation
+- Error handling
+
+Updated: 2025-01-21
+"""
+
+import tempfile
+import time
+from pathlib import Path
+
+import pytest
+
+from atlas.modules.mcp_tools.token_storage import (
+    AuthenticationRequiredException,
+    MCPTokenStorage,
+    StoredToken,
+    _make_token_key,
+    _parse_token_key,
+    get_token_storage,
+)
+
+
+class TestTokenKeyFunctions:
+    """Test helper functions for token key management."""
+
+    def test_make_token_key_basic(self):
+        """Should create key from email and server name."""
+        key = _make_token_key("user@example.com", "my-server")
+        assert key == "user@example.com:my-server"
+
+    def test_make_token_key_normalizes_case(self):
+        """Should normalize email to lowercase."""
+        key = _make_token_key("User@Example.COM", "My-Server")
+        assert key == "user@example.com:My-Server"
+
+    def test_parse_token_key_basic(self):
+        """Should parse key into email and server name."""
+        email, server = _parse_token_key("user@example.com:my-server")
+        assert email == "user@example.com"
+        assert server == "my-server"
+
+    def test_parse_token_key_with_colons_in_server(self):
+        """Should handle server names with colons (only splits on first colon)."""
+        email, server = _parse_token_key("user@example.com:server:with:colons")
+        assert email == "user@example.com"
+        assert server == "server:with:colons"
+
+    def test_parse_token_key_invalid_format(self):
+        """Should raise ValueError for invalid key format."""
+        with pytest.raises(ValueError, match="Invalid token key format"):
+            _parse_token_key("no-colon-here")
+
+
+class TestStoredToken:
+    """Test StoredToken dataclass."""
+
+    def test_stored_token_not_expired_no_expiry(self):
+        """Token without expiry should never be expired."""
+        token = StoredToken(
+            token_type="api_key",
+            token_value="test-key",
+            user_email="user@example.com",
+            server_name="test-server",
+            created_at=time.time(),
+            expires_at=None,
+        )
+        assert token.is_expired() is False
+
+    def test_stored_token_not_expired_future_expiry(self):
+        """Token with future expiry should not be expired."""
+        token = StoredToken(
+            token_type="api_key",
+            token_value="test-key",
+            user_email="user@example.com",
+            server_name="test-server",
+            created_at=time.time(),
+            expires_at=time.time() + 3600,  # 1 hour in future
+        )
+        assert token.is_expired() is False
+
+    def test_stored_token_expired_past_expiry(self):
+        """Token with past expiry should be expired."""
+        token = StoredToken(
+            token_type="api_key",
+            token_value="test-key",
+            user_email="user@example.com",
+            server_name="test-server",
+            created_at=time.time() - 7200,
+            expires_at=time.time() - 3600,  # 1 hour in past
+        )
+        assert token.is_expired() is True
+
+    def test_stored_token_expired_within_buffer(self):
+        """Token expiring within buffer should be considered expired."""
+        token = StoredToken(
+            token_type="api_key",
+            token_value="test-key",
+            user_email="user@example.com",
+            server_name="test-server",
+            created_at=time.time(),
+            expires_at=time.time() + 30,  # 30 seconds in future
+        )
+        # With 60-second buffer, should be considered expired
+        assert token.is_expired(buffer_seconds=60) is True
+        # With 10-second buffer, should not be expired
+        assert token.is_expired(buffer_seconds=10) is False
+
+    def test_stored_token_time_until_expiry(self):
+        """Should return correct time until expiry."""
+        future_time = time.time() + 3600
+        token = StoredToken(
+            token_type="api_key",
+            token_value="test-key",
+            user_email="user@example.com",
+            server_name="test-server",
+            created_at=time.time(),
+            expires_at=future_time,
+        )
+        # Should be approximately 3600 seconds (allow small margin)
+        time_until = token.time_until_expiry()
+        assert time_until is not None
+        assert 3590 < time_until <= 3600
+
+    def test_stored_token_time_until_expiry_none(self):
+        """Should return None when no expiry set."""
+        token = StoredToken(
+            token_type="api_key",
+            token_value="test-key",
+            user_email="user@example.com",
+            server_name="test-server",
+            created_at=time.time(),
+            expires_at=None,
+        )
+        assert token.time_until_expiry() is None
+
+    def test_stored_token_time_until_expiry_past(self):
+        """Should return 0 when already expired."""
+        token = StoredToken(
+            token_type="api_key",
+            token_value="test-key",
+            user_email="user@example.com",
+            server_name="test-server",
+            created_at=time.time() - 7200,
+            expires_at=time.time() - 3600,
+        )
+        assert token.time_until_expiry() == 0
+
+
+class TestAuthenticationRequiredException:
+    """Test AuthenticationRequiredException."""
+
+    def test_exception_basic(self):
+        """Should create exception with required fields."""
+        exc = AuthenticationRequiredException(
+            server_name="my-server",
+            auth_type="api_key",
+        )
+        assert exc.server_name == "my-server"
+        assert exc.auth_type == "api_key"
+        assert exc.message == "Authentication required"
+        assert exc.oauth_start_url is None
+
+    def test_exception_with_message(self):
+        """Should accept custom message."""
+        exc = AuthenticationRequiredException(
+            server_name="my-server",
+            auth_type="jwt",
+            message="Custom auth message",
+        )
+        assert exc.message == "Custom auth message"
+        assert str(exc) == "Custom auth message"
+
+    def test_exception_with_oauth_url(self):
+        """Should store OAuth start URL for OAuth auth type."""
+        exc = AuthenticationRequiredException(
+            server_name="oauth-server",
+            auth_type="oauth",
+            oauth_start_url="/api/mcp/auth/oauth-server/oauth/start",
+        )
+        assert exc.oauth_start_url == "/api/mcp/auth/oauth-server/oauth/start"
+
+    def test_exception_to_dict(self):
+        """Should convert to dict for frontend consumption."""
+        exc = AuthenticationRequiredException(
+            server_name="my-server",
+            auth_type="api_key",
+            message="Please provide API key",
+            oauth_start_url=None,
+        )
+        result = exc.to_dict()
+        assert result == {
+            "server_name": "my-server",
+            "auth_type": "api_key",
+            "message": "Please provide API key",
+            "oauth_start_url": None,
+        }
+
+
+class TestMCPTokenStorage:
+    """Test MCPTokenStorage class."""
+
+    @pytest.fixture
+    def temp_storage_dir(self):
+        """Create a temporary directory for token storage."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield Path(tmpdir)
+
+    @pytest.fixture
+    def storage(self, temp_storage_dir):
+        """Create a MCPTokenStorage instance with temp directory."""
+        # Pass encryption key directly to avoid app_settings caching issues
+        storage = MCPTokenStorage(
+            storage_dir=temp_storage_dir,
+            encryption_key="test-encryption-key-12345"
+        )
+        yield storage
+
+    def test_store_and_retrieve_token(self, storage):
+        """Should store and retrieve a token successfully."""
+        # Store token
+        stored = storage.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="my-api-key-123",
+            token_type="api_key",
+        )
+        assert stored.token_value == "my-api-key-123"
+        assert stored.token_type == "api_key"
+
+        # Retrieve token
+        retrieved = storage.get_token("user@example.com", "test-server")
+        assert retrieved is not None
+        assert retrieved.token_value == "my-api-key-123"
+        assert retrieved.token_type == "api_key"
+        assert retrieved.user_email == "user@example.com"
+        assert retrieved.server_name == "test-server"
+
+    def test_store_token_with_expiry(self, storage):
+        """Should store token with expiration time."""
+        expiry = time.time() + 3600
+        stored = storage.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="expiring-token",
+            token_type="jwt",
+            expires_at=expiry,
+        )
+        assert stored.expires_at == expiry
+
+        retrieved = storage.get_token("user@example.com", "test-server")
+        assert retrieved is not None
+        assert retrieved.expires_at == expiry
+        assert retrieved.is_expired() is False
+
+    def test_store_token_with_scopes(self, storage):
+        """Should store token with scopes."""
+        stored = storage.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="scoped-token",
+            token_type="bearer",
+            scopes="read write admin",
+        )
+        assert stored.scopes == "read write admin"
+
+        retrieved = storage.get_token("user@example.com", "test-server")
+        assert retrieved is not None
+        assert retrieved.scopes == "read write admin"
+
+    def test_get_nonexistent_token(self, storage):
+        """Should return None for nonexistent token."""
+        retrieved = storage.get_token("nobody@example.com", "no-server")
+        assert retrieved is None
+
+    def test_remove_token(self, storage):
+        """Should remove token successfully."""
+        # Store token
+        storage.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="to-be-removed",
+            token_type="api_key",
+        )
+
+        # Verify it exists
+        assert storage.get_token("user@example.com", "test-server") is not None
+
+        # Remove it
+        removed = storage.remove_token("user@example.com", "test-server")
+        assert removed is True
+
+        # Verify it's gone
+        assert storage.get_token("user@example.com", "test-server") is None
+
+    def test_remove_nonexistent_token(self, storage):
+        """Should return False when removing nonexistent token."""
+        removed = storage.remove_token("nobody@example.com", "no-server")
+        assert removed is False
+
+    def test_user_isolation(self, storage):
+        """Tokens should be isolated per user."""
+        # Store token for user1
+        storage.store_token(
+            user_email="user1@example.com",
+            server_name="shared-server",
+            token_value="user1-token",
+            token_type="api_key",
+        )
+
+        # Store token for user2
+        storage.store_token(
+            user_email="user2@example.com",
+            server_name="shared-server",
+            token_value="user2-token",
+            token_type="api_key",
+        )
+
+        # Each user should only see their own token
+        user1_token = storage.get_token("user1@example.com", "shared-server")
+        user2_token = storage.get_token("user2@example.com", "shared-server")
+
+        assert user1_token is not None
+        assert user2_token is not None
+        assert user1_token.token_value == "user1-token"
+        assert user2_token.token_value == "user2-token"
+
+    def test_overwrite_existing_token(self, storage):
+        """Should overwrite existing token for same user/server."""
+        # Store initial token
+        storage.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="old-token",
+            token_type="api_key",
+        )
+
+        # Store new token
+        storage.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="new-token",
+            token_type="jwt",
+        )
+
+        # Should get the new token
+        retrieved = storage.get_token("user@example.com", "test-server")
+        assert retrieved is not None
+        assert retrieved.token_value == "new-token"
+        assert retrieved.token_type == "jwt"
+
+    def test_get_user_auth_status(self, storage):
+        """Should return auth status for all user's tokens."""
+        # Store tokens for user
+        storage.store_token(
+            user_email="user@example.com",
+            server_name="server1",
+            token_value="token1",
+            token_type="api_key",
+        )
+        storage.store_token(
+            user_email="user@example.com",
+            server_name="server2",
+            token_value="token2",
+            token_type="jwt",
+            expires_at=time.time() + 3600,
+        )
+
+        # Store token for different user (should not be included)
+        storage.store_token(
+            user_email="other@example.com",
+            server_name="server3",
+            token_value="other-token",
+            token_type="bearer",
+        )
+
+        status = storage.get_user_auth_status("user@example.com")
+
+        assert "server1" in status
+        assert "server2" in status
+        assert "server3" not in status  # Different user
+
+        assert status["server1"]["token_type"] == "api_key"
+        assert status["server2"]["token_type"] == "jwt"
+        assert status["server2"]["is_expired"] is False
+
+    def test_email_case_insensitive(self, storage):
+        """Email lookups should be case-insensitive."""
+        storage.store_token(
+            user_email="User@Example.COM",
+            server_name="test-server",
+            token_value="test-token",
+            token_type="api_key",
+        )
+
+        # Should find with different case
+        retrieved = storage.get_token("user@example.com", "test-server")
+        assert retrieved is not None
+        assert retrieved.token_value == "test-token"
+
+
+class TestMCPTokenStoragePersistence:
+    """Test token storage persistence across instances."""
+
+    @pytest.fixture
+    def temp_storage_dir(self):
+        """Create a temporary directory for token storage."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield Path(tmpdir)
+
+    def test_persistence_across_instances(self, temp_storage_dir):
+        """Tokens should persist across storage instances."""
+        # Pass encryption_key directly to avoid app_settings caching issues
+        encryption_key = "test-encryption-key-12345"
+
+        # Create first storage instance and store token
+        storage1 = MCPTokenStorage(storage_dir=temp_storage_dir, encryption_key=encryption_key)
+        storage1.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="persistent-token",
+            token_type="api_key",
+        )
+
+        # Create second storage instance and retrieve token
+        storage2 = MCPTokenStorage(storage_dir=temp_storage_dir, encryption_key=encryption_key)
+        retrieved = storage2.get_token("user@example.com", "test-server")
+
+        assert retrieved is not None
+        assert retrieved.token_value == "persistent-token"
+
+
+class TestMCPTokenStorageEncryption:
+    """Test token encryption functionality."""
+
+    @pytest.fixture
+    def temp_storage_dir(self):
+        """Create a temporary directory for token storage."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield Path(tmpdir)
+
+    def test_tokens_encrypted_at_rest(self, temp_storage_dir):
+        """Token values should be encrypted in storage file."""
+        # Pass encryption key directly to avoid app_settings caching issues
+        storage = MCPTokenStorage(
+            storage_dir=temp_storage_dir,
+            encryption_key="test-encryption-key-12345"
+        )
+        storage.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="secret-api-key-xyz",
+            token_type="api_key",
+        )
+
+        # Read raw storage file
+        storage_file = temp_storage_dir / "mcp_tokens.enc"
+        raw_content = storage_file.read_text()
+
+        # Plain token should not appear in raw content
+        assert "secret-api-key-xyz" not in raw_content
+
+    def test_different_keys_cannot_decrypt(self, temp_storage_dir):
+        """Tokens encrypted with different keys should not be readable."""
+        # Store with one key
+        storage1 = MCPTokenStorage(
+            storage_dir=temp_storage_dir,
+            encryption_key="first-encryption-key"
+        )
+        storage1.store_token(
+            user_email="user@example.com",
+            server_name="test-server",
+            token_value="secret-token",
+            token_type="api_key",
+        )
+
+        # Try to read with different key
+        storage2 = MCPTokenStorage(
+            storage_dir=temp_storage_dir,
+            encryption_key="different-encryption-key"
+        )
+
+        # Should return None (decryption fails gracefully)
+        retrieved = storage2.get_token("user@example.com", "test-server")
+        assert retrieved is None
+
+
+class TestGetMCPTokenStorageSingleton:
+    """Test the get_token_storage singleton function."""
+
+    def test_returns_token_storage_instance(self):
+        """Should return a MCPTokenStorage instance."""
+        storage = get_token_storage()
+        assert isinstance(storage, MCPTokenStorage)
+
+    def test_returns_same_instance(self):
+        """Should return the same instance on repeated calls."""
+        storage1 = get_token_storage()
+        storage2 = get_token_storage()
+        assert storage1 is storage2

atlas/tests/test_tool_approval_config.py

@@ -0,0 +1,93 @@
+"""Tests for tool approval configuration loading and management."""
+
+from atlas.modules.config.config_manager import ConfigManager
+
+
+class TestToolApprovalConfig:
+    """Test tool approval configuration loading."""
+
+    def test_tool_approvals_config_loads(self):
+        """Test that tool approvals config loads successfully."""
+        cm = ConfigManager()
+        approval_config = cm.tool_approvals_config
+
+        assert approval_config is not None
+        assert hasattr(approval_config, "require_approval_by_default")
+        assert hasattr(approval_config, "tools")
+
+    def test_default_approval_config_structure(self):
+        """Test the structure of default approval config."""
+        cm = ConfigManager()
+        approval_config = cm.tool_approvals_config
+
+        # Default config should have require_approval_by_default (check it's boolean)
+        assert isinstance(approval_config.require_approval_by_default, bool)
+        # Default config should have tools dict (may or may not be empty)
+        assert isinstance(approval_config.tools, dict)
+
+    def test_tool_specific_config(self):
+        """Test that tool-specific configurations can be loaded."""
+        cm = ConfigManager()
+        approval_config = cm.tool_approvals_config
+
+        # Test basic structure - config may have tool-specific configs from overrides
+        assert hasattr(approval_config, 'tools')
+        assert isinstance(approval_config.tools, dict)
+
+        # If there are any tool configs, verify they have the right structure
+        for tool_name, tool_config in approval_config.tools.items():
+            assert hasattr(tool_config, 'require_approval')
+            assert hasattr(tool_config, 'allow_edit')
+            assert isinstance(tool_config.require_approval, bool)
+            assert isinstance(tool_config.allow_edit, bool)
+
+    def test_config_has_boolean_default(self):
+        """Test that require_approval_by_default is a boolean."""
+        cm = ConfigManager()
+        approval_config = cm.tool_approvals_config
+
+        assert isinstance(approval_config.require_approval_by_default, bool)
+
+    def test_tools_config_structure(self):
+        """Test that tools in config have correct structure."""
+        cm = ConfigManager()
+        approval_config = cm.tool_approvals_config
+
+        # Each tool config should have require_approval and allow_edit
+        for tool_name, tool_config in approval_config.tools.items():
+            assert hasattr(tool_config, 'require_approval')
+            assert hasattr(tool_config, 'allow_edit')
+            assert isinstance(tool_config.require_approval, bool)
+            assert isinstance(tool_config.allow_edit, bool)
+
+    def test_config_manager_provides_approvals_config(self):
+        """Test that ConfigManager provides tool_approvals_config."""
+        cm = ConfigManager()
+
+        assert hasattr(cm, 'tool_approvals_config')
+        assert cm.tool_approvals_config is not None
+
+    def test_multiple_config_manager_instances(self):
+        """Test that multiple ConfigManager instances can coexist."""
+        cm1 = ConfigManager()
+        cm2 = ConfigManager()
+
+        config1 = cm1.tool_approvals_config
+        config2 = cm2.tool_approvals_config
+
+        # Both should have valid configs
+        assert config1 is not None
+        assert config2 is not None
+
+    def test_config_contains_expected_fields(self):
+        """Test that approval config has all expected fields."""
+        cm = ConfigManager()
+        approval_config = cm.tool_approvals_config
+
+        # Should have these attributes
+        assert hasattr(approval_config, 'require_approval_by_default')
+        assert hasattr(approval_config, 'tools')
+
+        # Types should be correct
+        assert isinstance(approval_config.require_approval_by_default, bool)
+        assert isinstance(approval_config.tools, dict)