Initial commit: MCP Hub Community Edition v3.0.0
Community edition generated from private repo via sync pipeline. Includes 9 plugins (WordPress, WooCommerce, WP Advanced, Gitea, n8n, Supabase, OpenPanel, Appwrite, Directus) with ~587 tools. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
103
core/__init__.py
Normal file
103
core/__init__.py
Normal file
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Core modules for MCP server
|
||||
|
||||
Architecture:
|
||||
- Multi-Endpoint: Separate MCP endpoints for different plugin types
|
||||
- Tool Registry: Central tool management
|
||||
- Site Manager: Multi-site configuration
|
||||
- Middleware: Authentication, rate limiting, audit logging
|
||||
"""
|
||||
|
||||
# Authentication and API Keys
|
||||
from core.api_keys import APIKeyManager, get_api_key_manager
|
||||
|
||||
# Logging and Audit
|
||||
from core.audit_log import AuditLogger, EventType, LogLevel, get_audit_logger
|
||||
from core.auth import AuthManager, get_auth_manager
|
||||
|
||||
# Context Management
|
||||
from core.context import clear_api_key_context, get_api_key_context, set_api_key_context
|
||||
|
||||
# Multi-Endpoint Architecture (Phase X)
|
||||
from core.endpoints import (
|
||||
EndpointConfig,
|
||||
EndpointRegistry,
|
||||
EndpointType,
|
||||
MCPEndpointFactory,
|
||||
)
|
||||
|
||||
# Health Monitoring
|
||||
from core.health import (
|
||||
AlertThreshold,
|
||||
HealthMetric,
|
||||
HealthMonitor,
|
||||
ProjectHealthStatus,
|
||||
SystemMetrics,
|
||||
get_health_monitor,
|
||||
initialize_health_monitor,
|
||||
)
|
||||
|
||||
# Project and Site Management
|
||||
from core.project_manager import ProjectManager, get_project_manager
|
||||
|
||||
# Rate Limiting
|
||||
from core.rate_limiter import RateLimitConfig, RateLimiter, get_rate_limiter
|
||||
from core.site_manager import SiteConfig, SiteManager, get_site_manager
|
||||
|
||||
# Legacy (kept for backward compatibility, will be removed in v2.0)
|
||||
from core.site_registry import SiteInfo, SiteRegistry, get_site_registry
|
||||
from core.tool_generator import ToolGenerator
|
||||
|
||||
# Tool Management (Option B architecture)
|
||||
from core.tool_registry import ToolDefinition, ToolRegistry, get_tool_registry
|
||||
from core.unified_tools import UnifiedToolGenerator
|
||||
|
||||
__all__ = [
|
||||
# Authentication
|
||||
"AuthManager",
|
||||
"get_auth_manager",
|
||||
"APIKeyManager",
|
||||
"get_api_key_manager",
|
||||
# Project/Site Management
|
||||
"ProjectManager",
|
||||
"get_project_manager",
|
||||
"SiteManager",
|
||||
"SiteConfig",
|
||||
"get_site_manager",
|
||||
# Legacy (deprecated)
|
||||
"SiteRegistry",
|
||||
"SiteInfo",
|
||||
"get_site_registry",
|
||||
"UnifiedToolGenerator",
|
||||
# Tool Management
|
||||
"ToolRegistry",
|
||||
"ToolDefinition",
|
||||
"get_tool_registry",
|
||||
"ToolGenerator",
|
||||
# Multi-Endpoint Architecture
|
||||
"EndpointConfig",
|
||||
"EndpointType",
|
||||
"MCPEndpointFactory",
|
||||
"EndpointRegistry",
|
||||
# Logging
|
||||
"AuditLogger",
|
||||
"get_audit_logger",
|
||||
"LogLevel",
|
||||
"EventType",
|
||||
# Context
|
||||
"set_api_key_context",
|
||||
"get_api_key_context",
|
||||
"clear_api_key_context",
|
||||
# Health
|
||||
"HealthMonitor",
|
||||
"HealthMetric",
|
||||
"SystemMetrics",
|
||||
"ProjectHealthStatus",
|
||||
"AlertThreshold",
|
||||
"get_health_monitor",
|
||||
"initialize_health_monitor",
|
||||
# Rate Limiting
|
||||
"RateLimiter",
|
||||
"get_rate_limiter",
|
||||
"RateLimitConfig",
|
||||
]
|
||||
499
core/api_keys.py
Normal file
499
core/api_keys.py
Normal file
@@ -0,0 +1,499 @@
|
||||
"""
|
||||
API Key Management System
|
||||
|
||||
Comprehensive per-project API key management with scopes, expiration,
|
||||
and audit trail.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import secrets
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Valid scope values
|
||||
VALID_SCOPES = ["read", "write", "admin"]
|
||||
|
||||
# Scope can be single ("read") or multiple space-separated ("read write admin")
|
||||
Scope = str
|
||||
|
||||
def validate_scope(scope: str) -> bool:
|
||||
"""
|
||||
Validate that scope contains only valid scope values.
|
||||
|
||||
Args:
|
||||
scope: Single scope ("read") or space-separated scopes ("read write admin")
|
||||
|
||||
Returns:
|
||||
True if all scopes are valid, False otherwise
|
||||
"""
|
||||
if not scope:
|
||||
return False
|
||||
|
||||
scope_list = scope.split()
|
||||
return all(s in VALID_SCOPES for s in scope_list)
|
||||
|
||||
def normalize_scope(scope: str) -> str:
|
||||
"""
|
||||
Normalize scope string by removing duplicates and sorting.
|
||||
|
||||
Args:
|
||||
scope: Single scope or space-separated scopes
|
||||
|
||||
Returns:
|
||||
Normalized scope string (e.g., "admin read write" -> "read write admin")
|
||||
"""
|
||||
scope_list = scope.split()
|
||||
# Remove duplicates, sort by priority (read < write < admin)
|
||||
unique_scopes = []
|
||||
for s in ["read", "write", "admin"]:
|
||||
if s in scope_list:
|
||||
unique_scopes.append(s)
|
||||
return " ".join(unique_scopes)
|
||||
|
||||
@dataclass
|
||||
class APIKey:
|
||||
"""
|
||||
Represents an API key with metadata.
|
||||
|
||||
Fields:
|
||||
key_id: Unique identifier for the key
|
||||
key_hash: SHA256 hash of the actual key (for storage)
|
||||
project_id: Project this key belongs to ("*" for all projects)
|
||||
scope: Access scope - single ("read") or multiple space-separated ("read write admin")
|
||||
created_at: ISO timestamp when created
|
||||
expires_at: Optional ISO timestamp when expires
|
||||
last_used_at: Optional ISO timestamp of last use
|
||||
usage_count: Number of times used
|
||||
description: Optional description
|
||||
revoked: Whether the key has been revoked
|
||||
"""
|
||||
|
||||
key_id: str
|
||||
key_hash: str
|
||||
project_id: str
|
||||
scope: Scope
|
||||
created_at: str
|
||||
expires_at: str | None = None
|
||||
last_used_at: str | None = None
|
||||
usage_count: int = 0
|
||||
description: str | None = None
|
||||
revoked: bool = False
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> "APIKey":
|
||||
"""Create from dictionary."""
|
||||
return cls(**data)
|
||||
|
||||
def is_expired(self) -> bool:
|
||||
"""Check if key has expired."""
|
||||
if not self.expires_at:
|
||||
return False
|
||||
expires = datetime.fromisoformat(self.expires_at)
|
||||
return datetime.now() > expires
|
||||
|
||||
def is_valid(self) -> bool:
|
||||
"""Check if key is valid (not revoked, not expired)."""
|
||||
return not self.revoked and not self.is_expired()
|
||||
|
||||
class APIKeyManager:
|
||||
"""
|
||||
Manages per-project API keys with persistence.
|
||||
|
||||
Features:
|
||||
- Persistent JSON storage
|
||||
- Key creation with scopes
|
||||
- Key validation with project and scope checking
|
||||
- Key rotation and revocation
|
||||
- Usage tracking
|
||||
- Expiration support
|
||||
"""
|
||||
|
||||
def __init__(self, storage_path: str = "data/api_keys.json"):
|
||||
"""
|
||||
Initialize API Key Manager.
|
||||
|
||||
Args:
|
||||
storage_path: Path to JSON file for key storage
|
||||
"""
|
||||
self.storage_path = Path(storage_path)
|
||||
self.keys: dict[str, APIKey] = {}
|
||||
|
||||
# Ensure storage directory exists (with graceful fallback)
|
||||
try:
|
||||
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
except PermissionError:
|
||||
# Fallback to /tmp if we can't create in data/
|
||||
logger.warning(
|
||||
f"Cannot create directory {self.storage_path.parent}, " f"falling back to /tmp"
|
||||
)
|
||||
self.storage_path = Path("/tmp/api_keys.json")
|
||||
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Load existing keys
|
||||
self._load_keys()
|
||||
|
||||
logger.info(
|
||||
f"API Key Manager initialized with {len(self.keys)} keys "
|
||||
f"(storage: {self.storage_path})"
|
||||
)
|
||||
|
||||
def _load_keys(self) -> None:
|
||||
"""Load keys from storage file."""
|
||||
if not self.storage_path.exists():
|
||||
logger.info("No existing keys file found, starting fresh")
|
||||
return
|
||||
|
||||
try:
|
||||
with open(self.storage_path) as f:
|
||||
data = json.load(f)
|
||||
self.keys = {
|
||||
key_id: APIKey.from_dict(key_data) for key_id, key_data in data.items()
|
||||
}
|
||||
logger.info(f"Loaded {len(self.keys)} keys from storage")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load keys: {e}")
|
||||
self.keys = {}
|
||||
|
||||
def _save_keys(self) -> None:
|
||||
"""Save keys to storage file."""
|
||||
try:
|
||||
data = {key_id: key.to_dict() for key_id, key in self.keys.items()}
|
||||
with open(self.storage_path, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
logger.debug(f"Saved {len(self.keys)} keys to storage")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save keys: {e}")
|
||||
|
||||
def _hash_key(self, api_key: str) -> str:
|
||||
"""Hash API key for storage."""
|
||||
return hashlib.sha256(api_key.encode()).hexdigest()
|
||||
|
||||
def create_key(
|
||||
self,
|
||||
project_id: str,
|
||||
scope: Scope = "read",
|
||||
expires_in_days: int | None = None,
|
||||
description: str | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
Create a new API key.
|
||||
|
||||
Args:
|
||||
project_id: Project ID ("*" for all projects)
|
||||
scope: Access scope - single ("read") or multiple space-separated ("read write admin")
|
||||
expires_in_days: Optional expiration in days
|
||||
description: Optional description
|
||||
|
||||
Returns:
|
||||
dict: {"key": actual_key, "key_id": key_id, "project_id": project_id, "scope": scope}
|
||||
|
||||
Raises:
|
||||
ValueError: If scope contains invalid values
|
||||
"""
|
||||
# Validate and normalize scope
|
||||
if not validate_scope(scope):
|
||||
raise ValueError(
|
||||
f"Invalid scope: {scope}. Must contain only: {', '.join(VALID_SCOPES)}"
|
||||
)
|
||||
normalized_scope = normalize_scope(scope)
|
||||
|
||||
# Generate secure random key
|
||||
api_key = f"cmp_{secrets.token_urlsafe(32)}"
|
||||
key_id = f"key_{secrets.token_urlsafe(16)}"
|
||||
key_hash = self._hash_key(api_key)
|
||||
|
||||
# Calculate expiration
|
||||
expires_at = None
|
||||
if expires_in_days:
|
||||
expires = datetime.now() + timedelta(days=expires_in_days)
|
||||
expires_at = expires.isoformat()
|
||||
|
||||
# Create key object
|
||||
key = APIKey(
|
||||
key_id=key_id,
|
||||
key_hash=key_hash,
|
||||
project_id=project_id,
|
||||
scope=normalized_scope,
|
||||
created_at=datetime.now().isoformat(),
|
||||
expires_at=expires_at,
|
||||
description=description,
|
||||
)
|
||||
|
||||
# Store and save
|
||||
self.keys[key_id] = key
|
||||
self._save_keys()
|
||||
|
||||
logger.info(
|
||||
f"Created API key {key_id} for project {project_id} " f"with scope '{normalized_scope}'"
|
||||
)
|
||||
|
||||
return {
|
||||
"key": api_key,
|
||||
"key_id": key_id,
|
||||
"scope": normalized_scope,
|
||||
"project_id": project_id,
|
||||
"expires_at": expires_at,
|
||||
}
|
||||
|
||||
def validate_key(
|
||||
self,
|
||||
api_key: str,
|
||||
project_id: str,
|
||||
required_scope: Scope = "read",
|
||||
skip_project_check: bool = False,
|
||||
) -> str | None:
|
||||
"""
|
||||
Validate API key for project and scope.
|
||||
|
||||
Args:
|
||||
api_key: The API key to validate
|
||||
project_id: Project to check access for
|
||||
required_scope: Minimum required scope
|
||||
skip_project_check: Skip project-level validation (for unified tools)
|
||||
|
||||
Returns:
|
||||
Optional[str]: key_id if valid, None otherwise
|
||||
"""
|
||||
key_hash = self._hash_key(api_key)
|
||||
|
||||
# Find key by hash
|
||||
for key_id, key in self.keys.items():
|
||||
if key.key_hash != key_hash:
|
||||
continue
|
||||
|
||||
# Check if valid (not revoked, not expired)
|
||||
if not key.is_valid():
|
||||
logger.warning(
|
||||
f"Key {key_id} is invalid "
|
||||
f"(revoked={key.revoked}, expired={key.is_expired()})"
|
||||
)
|
||||
return None
|
||||
|
||||
# Check project access (unless skipped for unified tools)
|
||||
if not skip_project_check:
|
||||
if key.project_id != "*" and key.project_id != project_id:
|
||||
logger.warning(f"Key {key_id} does not have access to project {project_id}")
|
||||
return None
|
||||
|
||||
# Check scope: key must have required_scope or higher
|
||||
# Scope hierarchy: admin > write > read
|
||||
scope_hierarchy = {"read": 0, "write": 1, "admin": 2}
|
||||
key_scopes = key.scope.split()
|
||||
|
||||
# Check if required_scope is directly present
|
||||
if required_scope in key_scopes:
|
||||
# Update usage tracking
|
||||
key.last_used_at = datetime.now().isoformat()
|
||||
key.usage_count += 1
|
||||
self._save_keys()
|
||||
|
||||
logger.debug(f"Key {key_id} validated successfully (scope: {key.scope})")
|
||||
return key_id
|
||||
|
||||
# Check if key has higher scope (e.g., admin covers write and read)
|
||||
key_level = max(scope_hierarchy.get(s, 0) for s in key_scopes)
|
||||
required_level = scope_hierarchy.get(required_scope, 0)
|
||||
|
||||
if key_level >= required_level:
|
||||
# Update usage tracking
|
||||
key.last_used_at = datetime.now().isoformat()
|
||||
key.usage_count += 1
|
||||
self._save_keys()
|
||||
|
||||
logger.debug(f"Key {key_id} validated successfully (scope: {key.scope})")
|
||||
return key_id
|
||||
|
||||
logger.warning(
|
||||
f"Key {key_id} has insufficient scope "
|
||||
f"({key.scope} does not include {required_scope})"
|
||||
)
|
||||
return None
|
||||
|
||||
logger.warning("No matching API key found")
|
||||
return None
|
||||
|
||||
def get_key_by_token(self, api_key: str) -> APIKey | None:
|
||||
"""
|
||||
Get API key object by token (without project validation).
|
||||
|
||||
This method looks up an API key by its raw token value and returns
|
||||
the APIKey object if found. Unlike validate_key(), it does not
|
||||
validate against a specific project or scope.
|
||||
|
||||
Args:
|
||||
api_key: The raw API key token (e.g., "cmp_xxx...")
|
||||
|
||||
Returns:
|
||||
Optional[APIKey]: The APIKey object if found, None otherwise
|
||||
"""
|
||||
key_hash = self._hash_key(api_key)
|
||||
|
||||
for key_id, key in self.keys.items():
|
||||
if key.key_hash == key_hash:
|
||||
logger.debug(f"Found API key {key_id} by token")
|
||||
return key
|
||||
|
||||
logger.debug("No API key found for provided token")
|
||||
return None
|
||||
|
||||
def revoke_key(self, key_id: str) -> bool:
|
||||
"""
|
||||
Revoke an API key.
|
||||
|
||||
Args:
|
||||
key_id: Key ID to revoke
|
||||
|
||||
Returns:
|
||||
bool: True if revoked successfully
|
||||
"""
|
||||
if key_id not in self.keys:
|
||||
logger.warning(f"Key {key_id} not found")
|
||||
return False
|
||||
|
||||
self.keys[key_id].revoked = True
|
||||
self._save_keys()
|
||||
|
||||
logger.info(f"Revoked API key {key_id}")
|
||||
return True
|
||||
|
||||
def delete_key(self, key_id: str) -> bool:
|
||||
"""
|
||||
Permanently delete an API key.
|
||||
|
||||
Args:
|
||||
key_id: Key ID to delete
|
||||
|
||||
Returns:
|
||||
bool: True if deleted successfully
|
||||
"""
|
||||
if key_id not in self.keys:
|
||||
logger.warning(f"Key {key_id} not found")
|
||||
return False
|
||||
|
||||
del self.keys[key_id]
|
||||
self._save_keys()
|
||||
|
||||
logger.info(f"Deleted API key {key_id}")
|
||||
return True
|
||||
|
||||
def list_keys(self, project_id: str | None = None, include_revoked: bool = False) -> list[dict]:
|
||||
"""
|
||||
List API keys.
|
||||
|
||||
Args:
|
||||
project_id: Optional filter by project
|
||||
include_revoked: Include revoked keys
|
||||
|
||||
Returns:
|
||||
List of key information (without actual keys)
|
||||
"""
|
||||
keys = []
|
||||
|
||||
for key_id, key in self.keys.items():
|
||||
# Filter by project
|
||||
if project_id and key.project_id != project_id and key.project_id != "*":
|
||||
continue
|
||||
|
||||
# Filter revoked
|
||||
if not include_revoked and key.revoked:
|
||||
continue
|
||||
|
||||
keys.append(
|
||||
{
|
||||
"key_id": key_id,
|
||||
"project_id": key.project_id,
|
||||
"scope": key.scope,
|
||||
"created_at": key.created_at,
|
||||
"expires_at": key.expires_at,
|
||||
"last_used_at": key.last_used_at,
|
||||
"usage_count": key.usage_count,
|
||||
"description": key.description,
|
||||
"revoked": key.revoked,
|
||||
"expired": key.is_expired(),
|
||||
"valid": key.is_valid(),
|
||||
}
|
||||
)
|
||||
|
||||
return keys
|
||||
|
||||
def rotate_keys(self, project_id: str) -> list[dict[str, str]]:
|
||||
"""
|
||||
Rotate all keys for a project.
|
||||
|
||||
Creates new keys with same scopes and revokes old ones.
|
||||
|
||||
Args:
|
||||
project_id: Project to rotate keys for
|
||||
|
||||
Returns:
|
||||
List of new key information
|
||||
"""
|
||||
old_keys = [
|
||||
key for key in self.keys.values() if key.project_id == project_id and key.is_valid()
|
||||
]
|
||||
|
||||
new_keys = []
|
||||
|
||||
for old_key in old_keys:
|
||||
# Create new key with same scope
|
||||
new_key_data = self.create_key(
|
||||
project_id=project_id,
|
||||
scope=old_key.scope,
|
||||
description=f"Rotated from {old_key.key_id}",
|
||||
)
|
||||
new_keys.append(new_key_data)
|
||||
|
||||
# Revoke old key
|
||||
self.revoke_key(old_key.key_id)
|
||||
|
||||
logger.info(f"Rotated {len(new_keys)} keys for project {project_id}")
|
||||
return new_keys
|
||||
|
||||
def get_key_info(self, key_id: str) -> dict | None:
|
||||
"""
|
||||
Get information about a specific key.
|
||||
|
||||
Args:
|
||||
key_id: Key ID
|
||||
|
||||
Returns:
|
||||
Key information or None
|
||||
"""
|
||||
if key_id not in self.keys:
|
||||
return None
|
||||
|
||||
key = self.keys[key_id]
|
||||
return {
|
||||
"key_id": key_id,
|
||||
"project_id": key.project_id,
|
||||
"scope": key.scope,
|
||||
"created_at": key.created_at,
|
||||
"expires_at": key.expires_at,
|
||||
"last_used_at": key.last_used_at,
|
||||
"usage_count": key.usage_count,
|
||||
"description": key.description,
|
||||
"revoked": key.revoked,
|
||||
"expired": key.is_expired(),
|
||||
"valid": key.is_valid(),
|
||||
}
|
||||
|
||||
# Global instance
|
||||
_api_key_manager: APIKeyManager | None = None
|
||||
|
||||
def get_api_key_manager() -> APIKeyManager:
|
||||
"""Get the global API Key Manager instance."""
|
||||
global _api_key_manager
|
||||
if _api_key_manager is None:
|
||||
storage_path = os.getenv("API_KEYS_STORAGE", "data/api_keys.json")
|
||||
_api_key_manager = APIKeyManager(storage_path)
|
||||
return _api_key_manager
|
||||
565
core/audit_log.py
Normal file
565
core/audit_log.py
Normal file
@@ -0,0 +1,565 @@
|
||||
"""
|
||||
Audit Logging System
|
||||
|
||||
Comprehensive audit logging for all MCP operations.
|
||||
Tracks tool calls, authentication attempts, and system events.
|
||||
|
||||
Features:
|
||||
- Structured JSON logging
|
||||
- Log rotation support
|
||||
- Query and filter capabilities
|
||||
- Export to JSON/CSV
|
||||
- GDPR-compliant (no sensitive data in logs)
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
class LogLevel(Enum):
|
||||
"""Log severity levels."""
|
||||
|
||||
INFO = "INFO"
|
||||
WARNING = "WARNING"
|
||||
ERROR = "ERROR"
|
||||
CRITICAL = "CRITICAL"
|
||||
|
||||
class EventType(Enum):
|
||||
"""Types of events to log."""
|
||||
|
||||
TOOL_CALL = "tool_call"
|
||||
AUTHENTICATION = "authentication"
|
||||
HEALTH_CHECK = "health_check"
|
||||
ERROR = "error"
|
||||
SYSTEM = "system"
|
||||
|
||||
class AuditLogger:
|
||||
"""
|
||||
Audit logging system for MCP operations.
|
||||
|
||||
Logs all important events to a structured JSON log file with
|
||||
rotation support and query capabilities.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
log_dir: str = "logs",
|
||||
log_file: str = "audit.log",
|
||||
max_file_size_mb: int = 10,
|
||||
backup_count: int = 5,
|
||||
):
|
||||
"""
|
||||
Initialize audit logger.
|
||||
|
||||
Args:
|
||||
log_dir: Directory for log files
|
||||
log_file: Log file name
|
||||
max_file_size_mb: Max size before rotation (MB)
|
||||
backup_count: Number of backup files to keep
|
||||
"""
|
||||
# Setup Python logger for internal logging
|
||||
self.logger = logging.getLogger("AuditLogger")
|
||||
|
||||
# Try to create log directory, fallback to /tmp if permission denied
|
||||
self.log_dir = Path(log_dir)
|
||||
|
||||
try:
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
except PermissionError:
|
||||
# Fallback to /tmp/logs for Docker containers
|
||||
self.logger.warning(f"Permission denied for {log_dir}, using /tmp/logs instead")
|
||||
self.log_dir = Path("/tmp/logs")
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to create log directory: {e}, using /tmp/logs")
|
||||
self.log_dir = Path("/tmp/logs")
|
||||
try:
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
except Exception as e2:
|
||||
self.logger.critical(f"Cannot create any log directory: {e2}, logging disabled")
|
||||
# Set to None to disable file logging
|
||||
self.log_dir = None
|
||||
|
||||
if self.log_dir:
|
||||
self.log_file = self.log_dir / log_file
|
||||
self.max_file_size = max_file_size_mb * 1024 * 1024 # Convert to bytes
|
||||
self.backup_count = backup_count
|
||||
self.logger.info(f"Audit logger initialized: {self.log_file}")
|
||||
else:
|
||||
self.log_file = None
|
||||
self.max_file_size = 0
|
||||
self.backup_count = 0
|
||||
self.logger.warning("Audit logging to file is disabled due to permission errors")
|
||||
|
||||
def _rotate_logs_if_needed(self) -> None:
|
||||
"""Rotate log files if size exceeds limit."""
|
||||
if not self.log_file or not self.log_file.exists():
|
||||
return
|
||||
|
||||
if self.log_file.stat().st_size >= self.max_file_size:
|
||||
# Rotate existing backup files
|
||||
for i in range(self.backup_count - 1, 0, -1):
|
||||
old_backup = self.log_dir / f"{self.log_file.name}.{i}"
|
||||
new_backup = self.log_dir / f"{self.log_file.name}.{i + 1}"
|
||||
|
||||
if old_backup.exists():
|
||||
if new_backup.exists():
|
||||
new_backup.unlink() # Delete oldest if at limit
|
||||
old_backup.rename(new_backup)
|
||||
|
||||
# Move current log to .1
|
||||
backup = self.log_dir / f"{self.log_file.name}.1"
|
||||
if backup.exists():
|
||||
backup.unlink()
|
||||
self.log_file.rename(backup)
|
||||
|
||||
self.logger.info(f"Log rotated: {self.log_file}")
|
||||
|
||||
def _write_log_entry(self, entry: dict[str, Any]) -> None:
|
||||
"""
|
||||
Write a log entry to file.
|
||||
|
||||
Args:
|
||||
entry: Log entry dictionary
|
||||
"""
|
||||
# Skip if logging is disabled
|
||||
if not self.log_file:
|
||||
return
|
||||
|
||||
self._rotate_logs_if_needed()
|
||||
|
||||
try:
|
||||
with open(self.log_file, "a", encoding="utf-8") as f:
|
||||
# Write as JSON line
|
||||
json.dump(entry, f, ensure_ascii=False)
|
||||
f.write("\n")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to write audit log: {e}", exc_info=True)
|
||||
|
||||
def log_tool_call(
|
||||
self,
|
||||
tool_name: str,
|
||||
site: str | None = None,
|
||||
project_id: str | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
result_summary: str | None = None,
|
||||
error: str | None = None,
|
||||
duration_ms: int | None = None,
|
||||
user_id: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Log a tool call.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool called
|
||||
site: Site ID or alias (for unified tools)
|
||||
project_id: Full project ID (for per-site tools)
|
||||
params: Tool parameters (sensitive data should be filtered)
|
||||
result_summary: Brief summary of result (not full response)
|
||||
error: Error message if failed
|
||||
duration_ms: Execution duration in milliseconds
|
||||
user_id: User identifier (if available)
|
||||
"""
|
||||
# Filter sensitive data from params
|
||||
safe_params = self._filter_sensitive_data(params) if params else None
|
||||
|
||||
entry = {
|
||||
"timestamp": datetime.now(UTC).isoformat(),
|
||||
"event_type": EventType.TOOL_CALL.value,
|
||||
"level": LogLevel.ERROR.value if error else LogLevel.INFO.value,
|
||||
"tool_name": tool_name,
|
||||
"site": site,
|
||||
"project_id": project_id,
|
||||
"params": safe_params,
|
||||
"result_summary": result_summary,
|
||||
"error": error,
|
||||
"duration_ms": duration_ms,
|
||||
"user_id": user_id,
|
||||
"success": error is None,
|
||||
}
|
||||
|
||||
self._write_log_entry(entry)
|
||||
|
||||
def log_authentication(
|
||||
self,
|
||||
success: bool,
|
||||
project_id: str | None = None,
|
||||
reason: str | None = None,
|
||||
ip_address: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Log an authentication attempt.
|
||||
|
||||
Args:
|
||||
success: Whether authentication succeeded
|
||||
project_id: Project being accessed (if known)
|
||||
reason: Failure reason if unsuccessful
|
||||
ip_address: Client IP address
|
||||
"""
|
||||
entry = {
|
||||
"timestamp": datetime.now(UTC).isoformat(),
|
||||
"event_type": EventType.AUTHENTICATION.value,
|
||||
"level": LogLevel.WARNING.value if not success else LogLevel.INFO.value,
|
||||
"success": success,
|
||||
"project_id": project_id,
|
||||
"reason": reason,
|
||||
"ip_address": ip_address,
|
||||
}
|
||||
|
||||
self._write_log_entry(entry)
|
||||
|
||||
def log_error(
|
||||
self,
|
||||
error_type: str,
|
||||
error_message: str,
|
||||
context: dict[str, Any] | None = None,
|
||||
stack_trace: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Log an error event.
|
||||
|
||||
Args:
|
||||
error_type: Type of error (e.g., 'ValidationError', 'APIError')
|
||||
error_message: Error message
|
||||
context: Additional context
|
||||
stack_trace: Stack trace if available
|
||||
"""
|
||||
entry = {
|
||||
"timestamp": datetime.now(UTC).isoformat(),
|
||||
"event_type": EventType.ERROR.value,
|
||||
"level": LogLevel.ERROR.value,
|
||||
"error_type": error_type,
|
||||
"error_message": error_message,
|
||||
"context": context,
|
||||
"stack_trace": stack_trace,
|
||||
}
|
||||
|
||||
self._write_log_entry(entry)
|
||||
|
||||
def log_system_event(
|
||||
self, event: str, details: dict[str, Any] | None = None, level: LogLevel = LogLevel.INFO
|
||||
) -> None:
|
||||
"""
|
||||
Log a system event.
|
||||
|
||||
Args:
|
||||
event: Event description
|
||||
details: Event details
|
||||
level: Log level
|
||||
"""
|
||||
entry = {
|
||||
"timestamp": datetime.now(UTC).isoformat(),
|
||||
"event_type": EventType.SYSTEM.value,
|
||||
"level": level.value,
|
||||
"event": event,
|
||||
"details": details,
|
||||
}
|
||||
|
||||
self._write_log_entry(entry)
|
||||
|
||||
def _filter_sensitive_data(self, data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Filter sensitive data from logs (GDPR compliance).
|
||||
|
||||
Args:
|
||||
data: Data to filter
|
||||
|
||||
Returns:
|
||||
Filtered data with sensitive fields masked
|
||||
"""
|
||||
if not data:
|
||||
return {}
|
||||
|
||||
sensitive_keys = {
|
||||
"password",
|
||||
"app_password",
|
||||
"token",
|
||||
"api_key",
|
||||
"secret",
|
||||
"credential",
|
||||
"auth",
|
||||
"private_key",
|
||||
"access_token",
|
||||
"refresh_token",
|
||||
}
|
||||
|
||||
filtered = {}
|
||||
for key, value in data.items():
|
||||
# Check if key contains sensitive words
|
||||
if any(sensitive in key.lower() for sensitive in sensitive_keys):
|
||||
filtered[key] = "[REDACTED]"
|
||||
elif isinstance(value, dict):
|
||||
filtered[key] = self._filter_sensitive_data(value)
|
||||
else:
|
||||
filtered[key] = value
|
||||
|
||||
return filtered
|
||||
|
||||
def get_logs(
|
||||
self,
|
||||
event_type: EventType | None = None,
|
||||
start_time: datetime | None = None,
|
||||
end_time: datetime | None = None,
|
||||
level: LogLevel | None = None,
|
||||
project_id: str | None = None,
|
||||
tool_name: str | None = None,
|
||||
success_only: bool | None = None,
|
||||
limit: int = 100,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Query audit logs with filters.
|
||||
|
||||
Args:
|
||||
event_type: Filter by event type
|
||||
start_time: Start of time range
|
||||
end_time: End of time range
|
||||
level: Filter by log level
|
||||
project_id: Filter by project
|
||||
tool_name: Filter by tool name
|
||||
success_only: Only successful operations
|
||||
limit: Maximum number of entries to return
|
||||
|
||||
Returns:
|
||||
List of log entries matching filters
|
||||
"""
|
||||
if not self.log_file or not self.log_file.exists():
|
||||
return []
|
||||
|
||||
results = []
|
||||
|
||||
try:
|
||||
with open(self.log_file, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
try:
|
||||
entry = json.loads(line)
|
||||
|
||||
# Apply filters
|
||||
if event_type and entry.get("event_type") != event_type.value:
|
||||
continue
|
||||
|
||||
if level and entry.get("level") != level.value:
|
||||
continue
|
||||
|
||||
if project_id and entry.get("project_id") != project_id:
|
||||
continue
|
||||
|
||||
if tool_name and entry.get("tool_name") != tool_name:
|
||||
continue
|
||||
|
||||
if success_only is not None:
|
||||
if entry.get("success") != success_only:
|
||||
continue
|
||||
|
||||
# Time range filter
|
||||
if start_time or end_time:
|
||||
entry_time = datetime.fromisoformat(entry.get("timestamp", ""))
|
||||
|
||||
if start_time and entry_time < start_time:
|
||||
continue
|
||||
|
||||
if end_time and entry_time > end_time:
|
||||
continue
|
||||
|
||||
results.append(entry)
|
||||
|
||||
if len(results) >= limit:
|
||||
break
|
||||
|
||||
except json.JSONDecodeError:
|
||||
self.logger.warning(f"Invalid JSON in log: {line[:50]}...")
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error reading logs: {e}", exc_info=True)
|
||||
|
||||
return results
|
||||
|
||||
def export_logs(self, output_path: str, format: str = "json", **filter_kwargs) -> bool:
|
||||
"""
|
||||
Export logs to a file.
|
||||
|
||||
Args:
|
||||
output_path: Output file path
|
||||
format: Export format ('json' or 'csv')
|
||||
**filter_kwargs: Filters to apply (same as get_logs)
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
logs = self.get_logs(**filter_kwargs)
|
||||
|
||||
try:
|
||||
if format == "json":
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
json.dump(logs, f, indent=2, ensure_ascii=False)
|
||||
|
||||
elif format == "csv":
|
||||
import csv
|
||||
|
||||
if not logs:
|
||||
return False
|
||||
|
||||
# Get all unique keys from logs
|
||||
keys = set()
|
||||
for log in logs:
|
||||
keys.update(log.keys())
|
||||
|
||||
with open(output_path, "w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=sorted(keys))
|
||||
writer.writeheader()
|
||||
writer.writerows(logs)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported format: {format}")
|
||||
|
||||
self.logger.info(f"Exported {len(logs)} logs to {output_path}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error exporting logs: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
def get_recent_entries(self, limit: int = 10) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Get the most recent log entries.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of entries to return
|
||||
|
||||
Returns:
|
||||
List of recent log entries (newest first)
|
||||
"""
|
||||
if not self.log_file or not self.log_file.exists():
|
||||
return []
|
||||
|
||||
entries = []
|
||||
|
||||
try:
|
||||
with open(self.log_file, encoding="utf-8") as f:
|
||||
# Read all lines and get the last N
|
||||
lines = f.readlines()
|
||||
|
||||
# Process lines in reverse order
|
||||
for line in reversed(lines):
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
try:
|
||||
entry = json.loads(line)
|
||||
|
||||
# Format the entry for display
|
||||
formatted_entry = {
|
||||
"timestamp": entry.get("timestamp", ""),
|
||||
"event_type": entry.get("event_type", "unknown"),
|
||||
"level": entry.get("level", "INFO"),
|
||||
"message": self._format_log_message(entry),
|
||||
"metadata": {
|
||||
"project_id": entry.get("project_id"),
|
||||
"tool_name": entry.get("tool_name"),
|
||||
"site": entry.get("site"),
|
||||
"duration_ms": entry.get("duration_ms"),
|
||||
"success": entry.get("success"),
|
||||
},
|
||||
}
|
||||
|
||||
entries.append(formatted_entry)
|
||||
|
||||
if len(entries) >= limit:
|
||||
break
|
||||
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error reading recent logs: {e}", exc_info=True)
|
||||
|
||||
return entries
|
||||
|
||||
def _format_log_message(self, entry: dict[str, Any]) -> str:
|
||||
"""Format a log entry into a human-readable message."""
|
||||
event_type = entry.get("event_type", "")
|
||||
|
||||
if event_type == EventType.TOOL_CALL.value:
|
||||
tool_name = entry.get("tool_name", "unknown")
|
||||
if entry.get("error"):
|
||||
return f"{tool_name} failed: {entry.get('error', '')[:50]}"
|
||||
return f"{tool_name}"
|
||||
|
||||
elif event_type == EventType.AUTHENTICATION.value:
|
||||
if entry.get("success"):
|
||||
return "Authentication successful"
|
||||
return f"Authentication failed: {entry.get('reason', 'unknown')}"
|
||||
|
||||
elif event_type == EventType.ERROR.value:
|
||||
return f"{entry.get('error_type', 'Error')}: {entry.get('error_message', '')[:50]}"
|
||||
|
||||
elif event_type == EventType.SYSTEM.value:
|
||||
return entry.get("event", "System event")
|
||||
|
||||
return entry.get("event", entry.get("message", "Unknown event"))
|
||||
|
||||
def get_statistics(self) -> dict[str, Any]:
|
||||
"""
|
||||
Get statistics about audit logs.
|
||||
|
||||
Returns:
|
||||
Dictionary with statistics
|
||||
"""
|
||||
all_logs = self.get_logs(limit=10000) # Get recent logs
|
||||
|
||||
if not all_logs:
|
||||
return {"total_entries": 0, "by_type": {}, "by_level": {}, "success_rate": 0.0}
|
||||
|
||||
# Count by type
|
||||
by_type = {}
|
||||
by_level = {}
|
||||
successful = 0
|
||||
total_with_success_field = 0
|
||||
|
||||
for entry in all_logs:
|
||||
# Count by type
|
||||
event_type = entry.get("event_type", "unknown")
|
||||
by_type[event_type] = by_type.get(event_type, 0) + 1
|
||||
|
||||
# Count by level
|
||||
level = entry.get("level", "unknown")
|
||||
by_level[level] = by_level.get(level, 0) + 1
|
||||
|
||||
# Success rate
|
||||
if "success" in entry:
|
||||
total_with_success_field += 1
|
||||
if entry["success"]:
|
||||
successful += 1
|
||||
|
||||
success_rate = (
|
||||
(successful / total_with_success_field * 100) if total_with_success_field > 0 else 0.0
|
||||
)
|
||||
|
||||
# Calculate log file size
|
||||
log_file_size_mb = 0
|
||||
if self.log_file and self.log_file.exists():
|
||||
log_file_size_mb = round(self.log_file.stat().st_size / (1024 * 1024), 2)
|
||||
|
||||
return {
|
||||
"total_entries": len(all_logs),
|
||||
"by_type": by_type,
|
||||
"by_level": by_level,
|
||||
"success_rate": round(success_rate, 2),
|
||||
"log_file_size_mb": log_file_size_mb,
|
||||
}
|
||||
|
||||
# Global audit logger instance
|
||||
_audit_logger: AuditLogger | None = None
|
||||
|
||||
def get_audit_logger() -> AuditLogger:
|
||||
"""Get the global audit logger instance."""
|
||||
global _audit_logger
|
||||
if _audit_logger is None:
|
||||
_audit_logger = AuditLogger()
|
||||
return _audit_logger
|
||||
126
core/auth.py
Normal file
126
core/auth.py
Normal file
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
Authentication system
|
||||
|
||||
Simple API key based authentication for MCP server.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import secrets
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class AuthManager:
|
||||
"""
|
||||
Manage authentication for MCP server.
|
||||
|
||||
Currently supports simple API key authentication.
|
||||
Future: Can extend to support per-project keys, JWT, etc.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize authentication manager."""
|
||||
# Load master API key from environment
|
||||
self.master_api_key = os.getenv("MASTER_API_KEY")
|
||||
|
||||
if not self.master_api_key:
|
||||
# Generate a random key if not provided (dev mode)
|
||||
self.master_api_key = secrets.token_urlsafe(32)
|
||||
logger.warning(
|
||||
"No MASTER_API_KEY environment variable found. "
|
||||
f"Generated temporary key: {self.master_api_key[:8]}***{self.master_api_key[-4:]} "
|
||||
"(set MASTER_API_KEY in .env for production use)"
|
||||
)
|
||||
|
||||
# Project-specific keys (future feature)
|
||||
self.project_keys = {}
|
||||
|
||||
logger.info("Authentication manager initialized")
|
||||
|
||||
def validate_master_key(self, api_key: str) -> bool:
|
||||
"""
|
||||
Validate master API key.
|
||||
|
||||
Args:
|
||||
api_key: API key to validate
|
||||
|
||||
Returns:
|
||||
bool: True if valid
|
||||
"""
|
||||
is_valid = secrets.compare_digest(api_key, self.master_api_key)
|
||||
|
||||
if not is_valid:
|
||||
logger.warning("Invalid API key attempt")
|
||||
|
||||
return is_valid
|
||||
|
||||
def validate_project_key(self, project_id: str, api_key: str) -> bool:
|
||||
"""
|
||||
Validate project-specific API key.
|
||||
|
||||
Args:
|
||||
project_id: Project identifier
|
||||
api_key: API key to validate
|
||||
|
||||
Returns:
|
||||
bool: True if valid
|
||||
"""
|
||||
if project_id not in self.project_keys:
|
||||
# No project-specific key, fall back to master key
|
||||
return self.validate_master_key(api_key)
|
||||
|
||||
project_key = self.project_keys[project_id]
|
||||
is_valid = secrets.compare_digest(api_key, project_key)
|
||||
|
||||
if not is_valid:
|
||||
logger.warning(f"Invalid project key for {project_id}")
|
||||
|
||||
return is_valid
|
||||
|
||||
def add_project_key(self, project_id: str, api_key: str | None = None) -> str:
|
||||
"""
|
||||
Add or generate a project-specific API key.
|
||||
|
||||
Args:
|
||||
project_id: Project identifier
|
||||
api_key: Optional pre-defined key, will generate if None
|
||||
|
||||
Returns:
|
||||
str: The API key (useful if generated)
|
||||
"""
|
||||
if api_key is None:
|
||||
api_key = secrets.token_urlsafe(32)
|
||||
|
||||
self.project_keys[project_id] = api_key
|
||||
logger.info(f"Added API key for project: {project_id}")
|
||||
|
||||
return api_key
|
||||
|
||||
def remove_project_key(self, project_id: str) -> None:
|
||||
"""
|
||||
Remove project-specific API key.
|
||||
|
||||
Args:
|
||||
project_id: Project identifier
|
||||
"""
|
||||
if project_id in self.project_keys:
|
||||
del self.project_keys[project_id]
|
||||
logger.info(f"Removed API key for project: {project_id}")
|
||||
|
||||
def get_master_key(self) -> str:
|
||||
"""Get the master API key (for display/setup purposes)."""
|
||||
return self.master_api_key
|
||||
|
||||
def has_project_key(self, project_id: str) -> bool:
|
||||
"""Check if a project has its own API key."""
|
||||
return project_id in self.project_keys
|
||||
|
||||
# Global authentication manager instance
|
||||
_auth_manager: AuthManager | None = None
|
||||
|
||||
def get_auth_manager() -> AuthManager:
|
||||
"""Get the global authentication manager instance."""
|
||||
global _auth_manager
|
||||
if _auth_manager is None:
|
||||
_auth_manager = AuthManager()
|
||||
return _auth_manager
|
||||
40
core/context.py
Normal file
40
core/context.py
Normal file
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Request Context Storage
|
||||
|
||||
Stores request-level information using contextvars for thread-safe access
|
||||
across async operations.
|
||||
"""
|
||||
|
||||
from contextvars import ContextVar
|
||||
from typing import Any
|
||||
|
||||
# Context variable for storing API key info during request processing
|
||||
# This allows unified handlers to check project access permissions
|
||||
_api_key_context: ContextVar[dict[str, Any] | None] = ContextVar("api_key_context", default=None)
|
||||
|
||||
def set_api_key_context(key_id: str, project_id: str, scope: str, is_global: bool) -> None:
|
||||
"""
|
||||
Store API key information in request context.
|
||||
|
||||
Args:
|
||||
key_id: API key identifier
|
||||
project_id: Project the key belongs to ('*' for global)
|
||||
scope: Access scope (read/write/admin)
|
||||
is_global: Whether this is a global key
|
||||
"""
|
||||
_api_key_context.set(
|
||||
{"key_id": key_id, "project_id": project_id, "scope": scope, "is_global": is_global}
|
||||
)
|
||||
|
||||
def get_api_key_context() -> dict[str, Any] | None:
|
||||
"""
|
||||
Retrieve API key information from request context.
|
||||
|
||||
Returns:
|
||||
Dict with key_id, project_id, scope, is_global or None
|
||||
"""
|
||||
return _api_key_context.get()
|
||||
|
||||
def clear_api_key_context() -> None:
|
||||
"""Clear API key context (for cleanup)."""
|
||||
_api_key_context.set(None)
|
||||
75
core/dashboard/__init__.py
Normal file
75
core/dashboard/__init__.py
Normal file
@@ -0,0 +1,75 @@
|
||||
"""
|
||||
Dashboard module for MCP Hub Web UI.
|
||||
|
||||
Phase K: Web UI Dashboard
|
||||
"""
|
||||
|
||||
from .auth import DashboardAuth, get_dashboard_auth
|
||||
from .routes import (
|
||||
dashboard_api_audit_logs,
|
||||
dashboard_api_health,
|
||||
dashboard_api_keys_create,
|
||||
dashboard_api_keys_delete,
|
||||
# K.3: API Keys routes
|
||||
dashboard_api_keys_list,
|
||||
dashboard_api_keys_revoke,
|
||||
dashboard_api_project_detail,
|
||||
dashboard_api_projects,
|
||||
dashboard_api_stats,
|
||||
# K.4: Audit Logs routes
|
||||
dashboard_audit_logs_list,
|
||||
# K.5: Health Monitoring routes
|
||||
dashboard_health_page,
|
||||
dashboard_health_projects_partial,
|
||||
dashboard_home,
|
||||
# K.1: Core routes
|
||||
dashboard_login_page,
|
||||
dashboard_login_submit,
|
||||
dashboard_logout,
|
||||
dashboard_oauth_clients_create,
|
||||
dashboard_oauth_clients_delete,
|
||||
# K.4: OAuth Clients routes
|
||||
dashboard_oauth_clients_list,
|
||||
dashboard_project_detail,
|
||||
dashboard_project_health_check,
|
||||
# K.2: Projects routes
|
||||
dashboard_projects_list,
|
||||
# K.5: Settings routes
|
||||
dashboard_settings_page,
|
||||
register_dashboard_routes,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DashboardAuth",
|
||||
"get_dashboard_auth",
|
||||
"register_dashboard_routes",
|
||||
# K.1
|
||||
"dashboard_login_page",
|
||||
"dashboard_login_submit",
|
||||
"dashboard_logout",
|
||||
"dashboard_home",
|
||||
"dashboard_api_stats",
|
||||
# K.2
|
||||
"dashboard_projects_list",
|
||||
"dashboard_project_detail",
|
||||
"dashboard_api_projects",
|
||||
"dashboard_api_project_detail",
|
||||
"dashboard_project_health_check",
|
||||
# K.3
|
||||
"dashboard_api_keys_list",
|
||||
"dashboard_api_keys_create",
|
||||
"dashboard_api_keys_revoke",
|
||||
"dashboard_api_keys_delete",
|
||||
# K.4 OAuth
|
||||
"dashboard_oauth_clients_list",
|
||||
"dashboard_oauth_clients_create",
|
||||
"dashboard_oauth_clients_delete",
|
||||
# K.4 Audit
|
||||
"dashboard_audit_logs_list",
|
||||
"dashboard_api_audit_logs",
|
||||
# K.5
|
||||
"dashboard_health_page",
|
||||
"dashboard_api_health",
|
||||
"dashboard_health_projects_partial",
|
||||
"dashboard_settings_page",
|
||||
]
|
||||
278
core/dashboard/auth.py
Normal file
278
core/dashboard/auth.py
Normal file
@@ -0,0 +1,278 @@
|
||||
"""
|
||||
Dashboard Authentication - Session-based authentication for Web UI.
|
||||
|
||||
Phase K.1: Core Infrastructure
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import secrets
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Optional
|
||||
|
||||
import jwt
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import RedirectResponse, Response
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Singleton instance
|
||||
_dashboard_auth: Optional["DashboardAuth"] = None
|
||||
|
||||
@dataclass
|
||||
class DashboardSession:
|
||||
"""Dashboard session information."""
|
||||
|
||||
session_id: str
|
||||
created_at: datetime
|
||||
expires_at: datetime
|
||||
user_type: str # "master" or "api_key"
|
||||
key_id: str | None = None # For API key sessions
|
||||
|
||||
class DashboardAuth:
|
||||
"""
|
||||
Dashboard authentication manager.
|
||||
|
||||
Handles session-based authentication using JWT tokens stored in httpOnly cookies.
|
||||
"""
|
||||
|
||||
COOKIE_NAME = "mcp_dashboard_session"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
secret_key: str | None = None,
|
||||
session_expiry_hours: int = 24,
|
||||
master_api_key: str | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize dashboard authentication.
|
||||
|
||||
Args:
|
||||
secret_key: Secret for JWT signing. Generated if not provided.
|
||||
session_expiry_hours: Session expiration in hours.
|
||||
master_api_key: Master API key for validation.
|
||||
"""
|
||||
self.secret_key = secret_key or os.environ.get(
|
||||
"DASHBOARD_SESSION_SECRET", os.environ.get("OAUTH_JWT_SECRET_KEY")
|
||||
)
|
||||
if not self.secret_key:
|
||||
self.secret_key = secrets.token_hex(32)
|
||||
logger.warning(
|
||||
"DASHBOARD_SESSION_SECRET not set. Generated random session secret. "
|
||||
"All dashboard sessions will be invalidated on restart. "
|
||||
"Set DASHBOARD_SESSION_SECRET in your .env for persistent sessions."
|
||||
)
|
||||
self.session_expiry_hours = int(
|
||||
os.environ.get("DASHBOARD_SESSION_EXPIRY_HOURS", session_expiry_hours)
|
||||
)
|
||||
self.master_api_key = master_api_key or os.environ.get("MASTER_API_KEY")
|
||||
|
||||
# Rate limiting for login attempts
|
||||
self._login_attempts: dict[str, list] = {} # IP -> list of timestamps
|
||||
self.max_login_attempts = int(os.environ.get("DASHBOARD_LOGIN_RATE_LIMIT", 5))
|
||||
|
||||
logger.info(f"DashboardAuth initialized with {self.session_expiry_hours}h session expiry")
|
||||
|
||||
def validate_api_key(self, api_key: str) -> tuple[bool, str, str | None]:
|
||||
"""
|
||||
Validate an API key for dashboard login.
|
||||
|
||||
Args:
|
||||
api_key: The API key to validate.
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, user_type, key_id)
|
||||
- user_type: "master" or "api_key"
|
||||
- key_id: Key ID for API keys, None for master
|
||||
"""
|
||||
if not api_key:
|
||||
return False, "", None
|
||||
|
||||
# Check master API key
|
||||
if self.master_api_key and secrets.compare_digest(api_key, self.master_api_key):
|
||||
return True, "master", None
|
||||
|
||||
# Check project API keys with admin scope
|
||||
try:
|
||||
from core.api_keys import get_api_key_manager
|
||||
|
||||
api_key_manager = get_api_key_manager()
|
||||
|
||||
key_info = api_key_manager.validate_key(api_key)
|
||||
if key_info and "admin" in key_info.get("scope", "").split():
|
||||
return True, "api_key", key_info.get("key_id")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error checking API key: {e}")
|
||||
|
||||
return False, "", None
|
||||
|
||||
def check_rate_limit(self, client_ip: str) -> bool:
|
||||
"""
|
||||
Check if login attempts are within rate limit.
|
||||
|
||||
Args:
|
||||
client_ip: Client IP address.
|
||||
|
||||
Returns:
|
||||
True if within limit, False if exceeded.
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
window = timedelta(minutes=1)
|
||||
|
||||
# Clean old attempts
|
||||
if client_ip in self._login_attempts:
|
||||
self._login_attempts[client_ip] = [
|
||||
ts for ts in self._login_attempts[client_ip] if now - ts < window
|
||||
]
|
||||
else:
|
||||
self._login_attempts[client_ip] = []
|
||||
|
||||
return len(self._login_attempts[client_ip]) < self.max_login_attempts
|
||||
|
||||
def record_login_attempt(self, client_ip: str):
|
||||
"""Record a login attempt for rate limiting."""
|
||||
if client_ip not in self._login_attempts:
|
||||
self._login_attempts[client_ip] = []
|
||||
self._login_attempts[client_ip].append(datetime.now(UTC))
|
||||
|
||||
def create_session(self, user_type: str, key_id: str | None = None) -> str:
|
||||
"""
|
||||
Create a new dashboard session.
|
||||
|
||||
Args:
|
||||
user_type: Type of user ("master" or "api_key").
|
||||
key_id: Key ID for API key sessions.
|
||||
|
||||
Returns:
|
||||
JWT session token.
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
expires_at = now + timedelta(hours=self.session_expiry_hours)
|
||||
session_id = secrets.token_hex(16)
|
||||
|
||||
payload = {
|
||||
"sid": session_id,
|
||||
"type": user_type,
|
||||
"iat": now.timestamp(),
|
||||
"exp": expires_at.timestamp(),
|
||||
}
|
||||
if key_id:
|
||||
payload["kid"] = key_id
|
||||
|
||||
token = jwt.encode(payload, self.secret_key, algorithm="HS256")
|
||||
logger.info(f"Dashboard session created: type={user_type}, expires={expires_at}")
|
||||
|
||||
return token
|
||||
|
||||
def validate_session(self, token: str) -> DashboardSession | None:
|
||||
"""
|
||||
Validate a session token.
|
||||
|
||||
Args:
|
||||
token: JWT session token.
|
||||
|
||||
Returns:
|
||||
DashboardSession if valid, None otherwise.
|
||||
"""
|
||||
if not token:
|
||||
return None
|
||||
|
||||
try:
|
||||
payload = jwt.decode(token, self.secret_key, algorithms=["HS256"])
|
||||
|
||||
return DashboardSession(
|
||||
session_id=payload["sid"],
|
||||
created_at=datetime.fromtimestamp(payload["iat"]),
|
||||
expires_at=datetime.fromtimestamp(payload["exp"]),
|
||||
user_type=payload["type"],
|
||||
key_id=payload.get("kid"),
|
||||
)
|
||||
except jwt.ExpiredSignatureError:
|
||||
logger.debug("Dashboard session expired")
|
||||
return None
|
||||
except jwt.InvalidTokenError as e:
|
||||
logger.warning(f"Invalid dashboard session token: {e}")
|
||||
return None
|
||||
|
||||
def get_session_from_request(self, request: Request) -> DashboardSession | None:
|
||||
"""
|
||||
Extract and validate session from request.
|
||||
|
||||
Args:
|
||||
request: Starlette request object.
|
||||
|
||||
Returns:
|
||||
DashboardSession if valid session exists, None otherwise.
|
||||
"""
|
||||
token = request.cookies.get(self.COOKIE_NAME)
|
||||
if not token:
|
||||
return None
|
||||
return self.validate_session(token)
|
||||
|
||||
def set_session_cookie(self, response: Response, token: str) -> Response:
|
||||
"""
|
||||
Set session cookie on response.
|
||||
|
||||
Args:
|
||||
response: Response object to modify.
|
||||
token: Session token to set.
|
||||
|
||||
Returns:
|
||||
Modified response.
|
||||
"""
|
||||
response.set_cookie(
|
||||
key=self.COOKIE_NAME,
|
||||
value=token,
|
||||
max_age=self.session_expiry_hours * 3600,
|
||||
httponly=True,
|
||||
secure=os.environ.get("DASHBOARD_SECURE_COOKIE", "true").lower() == "true",
|
||||
samesite="lax",
|
||||
path="/", # Allow cookie for both /dashboard and /api/dashboard
|
||||
)
|
||||
return response
|
||||
|
||||
def clear_session_cookie(self, response: Response) -> Response:
|
||||
"""
|
||||
Clear session cookie on response.
|
||||
|
||||
Args:
|
||||
response: Response object to modify.
|
||||
|
||||
Returns:
|
||||
Modified response.
|
||||
"""
|
||||
response.delete_cookie(
|
||||
key=self.COOKIE_NAME,
|
||||
path="/", # Match the path used in set_cookie
|
||||
)
|
||||
return response
|
||||
|
||||
def require_auth(self, request: Request) -> RedirectResponse | None:
|
||||
"""
|
||||
Check if request is authenticated, redirect to login if not.
|
||||
|
||||
Args:
|
||||
request: Starlette request object.
|
||||
|
||||
Returns:
|
||||
RedirectResponse to login page if not authenticated, None if OK.
|
||||
"""
|
||||
session = self.get_session_from_request(request)
|
||||
if not session:
|
||||
# Store original URL for redirect after login
|
||||
next_url = str(request.url.path)
|
||||
if request.url.query:
|
||||
next_url += f"?{request.url.query}"
|
||||
return RedirectResponse(
|
||||
url=f"/dashboard/login?next={next_url}",
|
||||
status_code=303,
|
||||
)
|
||||
return None
|
||||
|
||||
def get_dashboard_auth() -> DashboardAuth:
|
||||
"""Get or create the singleton DashboardAuth instance."""
|
||||
global _dashboard_auth
|
||||
if _dashboard_auth is None:
|
||||
_dashboard_auth = DashboardAuth()
|
||||
return _dashboard_auth
|
||||
2075
core/dashboard/routes.py
Normal file
2075
core/dashboard/routes.py
Normal file
File diff suppressed because it is too large
Load Diff
63
core/endpoints/__init__.py
Normal file
63
core/endpoints/__init__.py
Normal file
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
Multi-Endpoint Architecture for MCP Hub
|
||||
|
||||
This module provides a factory pattern for creating scoped MCP endpoints.
|
||||
Each endpoint exposes only the tools relevant to its purpose.
|
||||
|
||||
Endpoints:
|
||||
/mcp - Admin endpoint (all tools, requires Master API Key)
|
||||
/mcp/wordpress - WordPress tools only (92 tools)
|
||||
/mcp/wordpress-advanced - WordPress Advanced tools (22 tools)
|
||||
/mcp/gitea - Gitea tools only (55 tools)
|
||||
/mcp/project/{id} - Project-specific tools
|
||||
|
||||
Benefits:
|
||||
- Better security: Users only see tools they can access
|
||||
- Optimized context: Smaller tool lists for AI assistants
|
||||
- Scalability: Easy to add new plugin endpoints
|
||||
- Clear separation of concerns
|
||||
"""
|
||||
|
||||
from .config import (
|
||||
ENDPOINT_CONFIGS,
|
||||
EndpointConfig,
|
||||
EndpointType,
|
||||
create_project_endpoint_config,
|
||||
get_endpoint_config,
|
||||
)
|
||||
from .factory import MCPEndpointFactory
|
||||
from .middleware import (
|
||||
AuthContext,
|
||||
EndpointAuditMiddleware,
|
||||
EndpointAuthMiddleware,
|
||||
EndpointRateLimitMiddleware,
|
||||
create_endpoint_middleware,
|
||||
)
|
||||
from .registry import (
|
||||
EndpointInfo,
|
||||
EndpointRegistry,
|
||||
get_endpoint_registry,
|
||||
initialize_endpoint_registry,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Config
|
||||
"EndpointConfig",
|
||||
"EndpointType",
|
||||
"ENDPOINT_CONFIGS",
|
||||
"get_endpoint_config",
|
||||
"create_project_endpoint_config",
|
||||
# Factory
|
||||
"MCPEndpointFactory",
|
||||
# Registry
|
||||
"EndpointRegistry",
|
||||
"EndpointInfo",
|
||||
"get_endpoint_registry",
|
||||
"initialize_endpoint_registry",
|
||||
# Middleware
|
||||
"EndpointAuthMiddleware",
|
||||
"EndpointRateLimitMiddleware",
|
||||
"EndpointAuditMiddleware",
|
||||
"create_endpoint_middleware",
|
||||
"AuthContext",
|
||||
]
|
||||
350
core/endpoints/config.py
Normal file
350
core/endpoints/config.py
Normal file
@@ -0,0 +1,350 @@
|
||||
"""
|
||||
Endpoint Configuration Module
|
||||
|
||||
Defines configurations for different MCP endpoints.
|
||||
Each endpoint has specific plugin types, scopes, and access requirements.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
class EndpointType(Enum):
|
||||
"""Types of MCP endpoints"""
|
||||
|
||||
ADMIN = "admin"
|
||||
SYSTEM = "system" # Phase X.3 - System tools only
|
||||
WORDPRESS = "wordpress"
|
||||
WOOCOMMERCE = "woocommerce"
|
||||
WORDPRESS_ADVANCED = "wordpress_advanced"
|
||||
GITEA = "gitea"
|
||||
N8N = "n8n"
|
||||
SUPABASE = "supabase" # Phase G
|
||||
OPENPANEL = "openpanel" # Phase H
|
||||
APPWRITE = "appwrite" # Phase I
|
||||
DIRECTUS = "directus" # Phase J
|
||||
PROJECT = "project" # Dynamic per-project endpoint
|
||||
CUSTOM = "custom"
|
||||
|
||||
@dataclass
|
||||
class EndpointConfig:
|
||||
"""
|
||||
Configuration for a single MCP endpoint.
|
||||
|
||||
Attributes:
|
||||
path: URL mount path for this endpoint (e.g., "/wordpress" → /wordpress/mcp)
|
||||
name: Display name for the endpoint
|
||||
description: Human-readable description
|
||||
endpoint_type: Type of endpoint (admin, wordpress, etc.)
|
||||
plugin_types: List of plugin types to include
|
||||
require_master_key: Whether Master API Key is required
|
||||
allowed_scopes: Allowed API key scopes (empty = all)
|
||||
tool_whitelist: Specific tools to include (None = all from plugins)
|
||||
tool_blacklist: Specific tools to exclude
|
||||
site_filter: Filter to specific site (for project endpoints)
|
||||
max_tools: Maximum number of tools (for safety)
|
||||
"""
|
||||
|
||||
path: str
|
||||
name: str
|
||||
description: str
|
||||
endpoint_type: EndpointType
|
||||
plugin_types: list[str] = field(default_factory=list)
|
||||
require_master_key: bool = False
|
||||
allowed_scopes: set[str] = field(default_factory=set)
|
||||
tool_whitelist: set[str] | None = None
|
||||
tool_blacklist: set[str] = field(default_factory=set)
|
||||
site_filter: str | None = None
|
||||
max_tools: int = 200
|
||||
|
||||
def __post_init__(self):
|
||||
"""Validate configuration after initialization"""
|
||||
if not self.path.startswith("/"):
|
||||
raise ValueError(f"Endpoint path must start with '/': {self.path}")
|
||||
|
||||
if self.tool_whitelist and self.tool_blacklist:
|
||||
overlap = self.tool_whitelist & self.tool_blacklist
|
||||
if overlap:
|
||||
raise ValueError(f"Tools cannot be in both whitelist and blacklist: {overlap}")
|
||||
|
||||
def allows_plugin(self, plugin_type: str) -> bool:
|
||||
"""Check if this endpoint allows a specific plugin type"""
|
||||
if not self.plugin_types:
|
||||
return True # Empty list = all plugins
|
||||
return plugin_type in self.plugin_types
|
||||
|
||||
def allows_tool(self, tool_name: str) -> bool:
|
||||
"""Check if this endpoint allows a specific tool"""
|
||||
# Check blacklist first
|
||||
if tool_name in self.tool_blacklist:
|
||||
return False
|
||||
|
||||
# If whitelist exists, tool must be in it
|
||||
if self.tool_whitelist is not None:
|
||||
return tool_name in self.tool_whitelist
|
||||
|
||||
return True
|
||||
|
||||
def allows_scope(self, scope: str) -> bool:
|
||||
"""Check if this endpoint allows a specific API key scope"""
|
||||
if not self.allowed_scopes:
|
||||
return True # Empty set = all scopes
|
||||
return scope in self.allowed_scopes
|
||||
|
||||
# Predefined endpoint configurations
|
||||
ENDPOINT_CONFIGS = {
|
||||
# Admin endpoint - all tools, requires Master API Key
|
||||
# Mounted at "/" → /mcp (FastMCP adds /mcp automatically)
|
||||
EndpointType.ADMIN: EndpointConfig(
|
||||
path="/",
|
||||
name="Coolify Admin",
|
||||
description="Full administrative access to all tools and plugins",
|
||||
endpoint_type=EndpointType.ADMIN,
|
||||
plugin_types=[], # Empty = all plugins
|
||||
require_master_key=True,
|
||||
allowed_scopes={"admin"},
|
||||
max_tools=400,
|
||||
),
|
||||
# System endpoint - system tools only (17 tools) - Phase X.3
|
||||
# For API key management, OAuth, rate limiting without loading all plugins
|
||||
EndpointType.SYSTEM: EndpointConfig(
|
||||
path="/system",
|
||||
name="System Manager",
|
||||
description="System management tools (API keys, OAuth, health, rate limiting)",
|
||||
endpoint_type=EndpointType.SYSTEM,
|
||||
plugin_types=["system"], # Only system tools
|
||||
require_master_key=True,
|
||||
allowed_scopes={"admin"},
|
||||
# Whitelist only system tools
|
||||
tool_whitelist={
|
||||
# API Key Management (6)
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_list",
|
||||
"manage_api_keys_get_info",
|
||||
"manage_api_keys_revoke",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
# Health & Status (4)
|
||||
"list_projects",
|
||||
"get_endpoints",
|
||||
"get_system_info",
|
||||
"get_audit_log",
|
||||
# OAuth Management (4)
|
||||
"oauth_register_client",
|
||||
"oauth_list_clients",
|
||||
"oauth_revoke_client",
|
||||
"oauth_get_client_info",
|
||||
# Rate Limiting (3)
|
||||
"get_rate_limit_stats",
|
||||
"reset_rate_limit",
|
||||
"set_rate_limit_config",
|
||||
},
|
||||
max_tools=20,
|
||||
),
|
||||
# WordPress endpoint - core WordPress tools only (64 tools)
|
||||
EndpointType.WORDPRESS: EndpointConfig(
|
||||
path="/wordpress",
|
||||
name="WordPress Manager",
|
||||
description="WordPress content management tools (posts, pages, media, SEO, menus)",
|
||||
endpoint_type=EndpointType.WORDPRESS,
|
||||
plugin_types=["wordpress"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"read", "write", "admin"},
|
||||
# Blacklist system and admin tools
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
},
|
||||
max_tools=70,
|
||||
),
|
||||
# WooCommerce endpoint - e-commerce tools (28 tools)
|
||||
EndpointType.WOOCOMMERCE: EndpointConfig(
|
||||
path="/woocommerce",
|
||||
name="WooCommerce Manager",
|
||||
description="WooCommerce e-commerce tools (products, orders, customers, coupons, reports)",
|
||||
endpoint_type=EndpointType.WOOCOMMERCE,
|
||||
plugin_types=["woocommerce"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"read", "write", "admin"},
|
||||
# Blacklist system and admin tools
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
},
|
||||
max_tools=35,
|
||||
),
|
||||
# WordPress Advanced endpoint - advanced operations
|
||||
EndpointType.WORDPRESS_ADVANCED: EndpointConfig(
|
||||
path="/wordpress-advanced",
|
||||
name="WordPress Advanced",
|
||||
description="WordPress advanced operations (database, bulk, system)",
|
||||
endpoint_type=EndpointType.WORDPRESS_ADVANCED,
|
||||
plugin_types=["wordpress_advanced"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"admin"}, # Admin scope required
|
||||
max_tools=30,
|
||||
),
|
||||
# Gitea endpoint - Git repository management
|
||||
EndpointType.GITEA: EndpointConfig(
|
||||
path="/gitea",
|
||||
name="Gitea Manager",
|
||||
description="Git repository management tools (repos, issues, PRs)",
|
||||
endpoint_type=EndpointType.GITEA,
|
||||
plugin_types=["gitea"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"read", "write", "admin"},
|
||||
# Blacklist system and admin tools
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
},
|
||||
max_tools=60,
|
||||
),
|
||||
# n8n endpoint - Workflow automation management (60 tools)
|
||||
EndpointType.N8N: EndpointConfig(
|
||||
path="/n8n",
|
||||
name="n8n Automation",
|
||||
description="Workflow automation management (workflows, executions, credentials, tags)",
|
||||
endpoint_type=EndpointType.N8N,
|
||||
plugin_types=["n8n"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"read", "write", "admin"},
|
||||
# Blacklist system and admin tools
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
},
|
||||
max_tools=70,
|
||||
),
|
||||
# Supabase endpoint - Self-Hosted management (70 tools) - Phase G
|
||||
EndpointType.SUPABASE: EndpointConfig(
|
||||
path="/supabase",
|
||||
name="Supabase Manager",
|
||||
description="Supabase Self-Hosted management (database, auth, storage, functions, admin)",
|
||||
endpoint_type=EndpointType.SUPABASE,
|
||||
plugin_types=["supabase"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"read", "write", "admin"},
|
||||
# Blacklist system and admin tools
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
},
|
||||
max_tools=80,
|
||||
),
|
||||
# OpenPanel endpoint - Product Analytics (73 tools) - Phase H
|
||||
EndpointType.OPENPANEL: EndpointConfig(
|
||||
path="/openpanel",
|
||||
name="OpenPanel Analytics",
|
||||
description="OpenPanel product analytics management (events, export, funnels, dashboards)",
|
||||
endpoint_type=EndpointType.OPENPANEL,
|
||||
plugin_types=["openpanel"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"read", "write", "admin"},
|
||||
# Blacklist system and admin tools
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
},
|
||||
max_tools=80,
|
||||
),
|
||||
# Appwrite endpoint - Backend-as-a-Service (100 tools) - Phase I
|
||||
EndpointType.APPWRITE: EndpointConfig(
|
||||
path="/appwrite",
|
||||
name="Appwrite Manager",
|
||||
description="Appwrite Self-Hosted management (databases, documents, users, teams, storage, functions, messaging)",
|
||||
endpoint_type=EndpointType.APPWRITE,
|
||||
plugin_types=["appwrite"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"read", "write", "admin"},
|
||||
# Blacklist system and admin tools
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
},
|
||||
max_tools=110,
|
||||
),
|
||||
# Directus endpoint - Headless CMS (100 tools) - Phase J
|
||||
EndpointType.DIRECTUS: EndpointConfig(
|
||||
path="/directus",
|
||||
name="Directus CMS",
|
||||
description="Directus Self-Hosted CMS management (items, collections, files, users, roles, flows, dashboards)",
|
||||
endpoint_type=EndpointType.DIRECTUS,
|
||||
plugin_types=["directus"],
|
||||
require_master_key=False,
|
||||
allowed_scopes={"read", "write", "admin"},
|
||||
# Blacklist system and admin tools
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
},
|
||||
max_tools=110,
|
||||
),
|
||||
}
|
||||
|
||||
def get_endpoint_config(endpoint_type: EndpointType) -> EndpointConfig:
|
||||
"""Get configuration for a specific endpoint type"""
|
||||
if endpoint_type not in ENDPOINT_CONFIGS:
|
||||
raise ValueError(f"Unknown endpoint type: {endpoint_type}")
|
||||
return ENDPOINT_CONFIGS[endpoint_type]
|
||||
|
||||
def create_project_endpoint_config(
|
||||
project_id: str, plugin_type: str, site_alias: str | None = None
|
||||
) -> EndpointConfig:
|
||||
"""
|
||||
Create a dynamic endpoint configuration for a specific project.
|
||||
|
||||
Args:
|
||||
project_id: Full project ID (e.g., "wordpress_site4")
|
||||
plugin_type: Plugin type (e.g., "wordpress")
|
||||
site_alias: Optional site alias for the path
|
||||
|
||||
Returns:
|
||||
EndpointConfig for the project-specific endpoint
|
||||
"""
|
||||
path_suffix = site_alias or project_id
|
||||
|
||||
# FastMCP adds /mcp automatically, so /project/xxx → /project/xxx/mcp
|
||||
return EndpointConfig(
|
||||
path=f"/project/{path_suffix}",
|
||||
name=f"Project: {project_id}",
|
||||
description=f"Tools for project {project_id}",
|
||||
endpoint_type=EndpointType.PROJECT,
|
||||
plugin_types=[plugin_type],
|
||||
require_master_key=False,
|
||||
site_filter=project_id,
|
||||
# Blacklist admin tools for project endpoints
|
||||
tool_blacklist={
|
||||
"manage_api_keys_create",
|
||||
"manage_api_keys_delete",
|
||||
"manage_api_keys_rotate",
|
||||
"oauth_register_client",
|
||||
"oauth_revoke_client",
|
||||
"list_projects", # Only show own project
|
||||
"oauth_list_clients",
|
||||
},
|
||||
max_tools=120,
|
||||
)
|
||||
295
core/endpoints/factory.py
Normal file
295
core/endpoints/factory.py
Normal file
@@ -0,0 +1,295 @@
|
||||
"""
|
||||
MCP Endpoint Factory
|
||||
|
||||
Creates and configures FastMCP instances for different endpoints.
|
||||
Each endpoint has its own set of tools based on configuration.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.server.middleware import Middleware
|
||||
|
||||
from .config import EndpointConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from core.site_manager import SiteManager
|
||||
from core.tool_registry import ToolRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class MCPEndpointFactory:
|
||||
"""
|
||||
Factory for creating scoped MCP endpoints.
|
||||
|
||||
Each endpoint is a separate FastMCP instance with only
|
||||
the tools relevant to its configuration.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
site_manager: "SiteManager",
|
||||
tool_registry: "ToolRegistry",
|
||||
middleware_classes: list[type] | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the endpoint factory.
|
||||
|
||||
Args:
|
||||
site_manager: Site manager for accessing site configurations
|
||||
tool_registry: Central tool registry
|
||||
middleware_classes: List of middleware classes to apply to endpoints
|
||||
"""
|
||||
self.site_manager = site_manager
|
||||
self.tool_registry = tool_registry
|
||||
self.middleware_classes = middleware_classes or []
|
||||
self.endpoints: dict[str, FastMCP] = {}
|
||||
self._tool_handlers: dict[str, Callable] = {}
|
||||
|
||||
def register_tool_handler(self, tool_name: str, handler: Callable):
|
||||
"""
|
||||
Register a tool handler function.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
handler: Async handler function for the tool
|
||||
"""
|
||||
self._tool_handlers[tool_name] = handler
|
||||
|
||||
def create_endpoint(
|
||||
self, config: EndpointConfig, custom_middleware: list[Middleware] | None = None
|
||||
) -> FastMCP:
|
||||
"""
|
||||
Create a new MCP endpoint with scoped tools.
|
||||
|
||||
Args:
|
||||
config: Endpoint configuration
|
||||
custom_middleware: Additional middleware for this endpoint
|
||||
|
||||
Returns:
|
||||
Configured FastMCP instance
|
||||
"""
|
||||
logger.info(f"Creating endpoint: {config.path} ({config.name})")
|
||||
|
||||
# Create FastMCP instance
|
||||
mcp = FastMCP(config.name)
|
||||
|
||||
# Get tools for this endpoint
|
||||
tools = self._get_tools_for_endpoint(config)
|
||||
|
||||
logger.info(f" - Registering {len(tools)} tools for {config.path}")
|
||||
|
||||
# Register tools
|
||||
for tool_info in tools:
|
||||
self._register_tool(mcp, tool_info, config)
|
||||
|
||||
# Apply middleware
|
||||
if custom_middleware:
|
||||
for middleware in custom_middleware:
|
||||
mcp.add_middleware(middleware)
|
||||
|
||||
# Store endpoint
|
||||
self.endpoints[config.path] = mcp
|
||||
|
||||
logger.info(f" - Endpoint {config.path} created successfully")
|
||||
|
||||
return mcp
|
||||
|
||||
def _get_tools_for_endpoint(self, config: EndpointConfig) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Get list of tools that should be available on this endpoint.
|
||||
|
||||
Args:
|
||||
config: Endpoint configuration
|
||||
|
||||
Returns:
|
||||
List of tool definitions
|
||||
"""
|
||||
tools = []
|
||||
|
||||
# Get all tools from registry
|
||||
all_tools = self.tool_registry.get_all()
|
||||
|
||||
for tool_def in all_tools:
|
||||
tool_name = tool_def.name
|
||||
|
||||
# Check plugin type filter
|
||||
plugin_type = self._extract_plugin_type(tool_name)
|
||||
if plugin_type and not config.allows_plugin(plugin_type):
|
||||
continue
|
||||
|
||||
# Check tool whitelist/blacklist
|
||||
if not config.allows_tool(tool_name):
|
||||
continue
|
||||
|
||||
# For project endpoints, filter by site
|
||||
if config.site_filter and plugin_type:
|
||||
# Tool should work with the specific site
|
||||
pass # Site filtering happens at execution time
|
||||
|
||||
tools.append(
|
||||
{
|
||||
"name": tool_name,
|
||||
"description": tool_def.description,
|
||||
"parameters": tool_def.parameters,
|
||||
"handler": tool_def.handler,
|
||||
"plugin_type": plugin_type,
|
||||
}
|
||||
)
|
||||
|
||||
# Check max tools limit
|
||||
if len(tools) > config.max_tools:
|
||||
logger.warning(
|
||||
f"Endpoint {config.path} has {len(tools)} tools, "
|
||||
f"exceeding max_tools={config.max_tools}"
|
||||
)
|
||||
|
||||
return tools
|
||||
|
||||
def _extract_plugin_type(self, tool_name: str) -> str | None:
|
||||
"""
|
||||
Extract plugin type from tool name.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
|
||||
Returns:
|
||||
Plugin type or None for system tools
|
||||
"""
|
||||
# Check for wordpress_advanced first (before wordpress_)
|
||||
# Tools are named: wordpress_advanced_wp_db_*, wordpress_advanced_bulk_*, wordpress_advanced_system_*
|
||||
if tool_name.startswith("wordpress_advanced_"):
|
||||
return "wordpress_advanced"
|
||||
|
||||
if tool_name.startswith("wordpress_"):
|
||||
return "wordpress"
|
||||
|
||||
elif tool_name.startswith("woocommerce_"):
|
||||
return "woocommerce"
|
||||
|
||||
elif tool_name.startswith("gitea_"):
|
||||
return "gitea"
|
||||
|
||||
elif tool_name.startswith("n8n_"):
|
||||
return "n8n"
|
||||
|
||||
elif tool_name.startswith("supabase_"):
|
||||
return "supabase"
|
||||
|
||||
elif tool_name.startswith("openpanel_"):
|
||||
return "openpanel"
|
||||
|
||||
elif tool_name.startswith("appwrite_"):
|
||||
return "appwrite"
|
||||
|
||||
elif tool_name.startswith("directus_"):
|
||||
return "directus"
|
||||
|
||||
# System tools have no plugin type
|
||||
return None
|
||||
|
||||
def _register_tool(self, mcp: FastMCP, tool_info: dict[str, Any], config: EndpointConfig):
|
||||
"""
|
||||
Register a single tool with the FastMCP instance.
|
||||
|
||||
Args:
|
||||
mcp: FastMCP instance
|
||||
tool_info: Tool definition
|
||||
config: Endpoint configuration
|
||||
"""
|
||||
tool_name = tool_info["name"]
|
||||
|
||||
# Get handler
|
||||
handler = tool_info.get("handler") or self._tool_handlers.get(tool_name)
|
||||
|
||||
if not handler:
|
||||
logger.warning(f"No handler found for tool: {tool_name}")
|
||||
return
|
||||
|
||||
# Wrap handler with endpoint context
|
||||
wrapped_handler = self._wrap_handler(handler, tool_name, config)
|
||||
|
||||
# Register with FastMCP
|
||||
# We need to create a function with the correct signature
|
||||
mcp.tool()(wrapped_handler)
|
||||
|
||||
def _wrap_handler(self, handler: Callable, tool_name: str, config: EndpointConfig) -> Callable:
|
||||
"""
|
||||
Wrap a tool handler with endpoint-specific logic.
|
||||
|
||||
Args:
|
||||
handler: Original handler function
|
||||
tool_name: Name of the tool
|
||||
config: Endpoint configuration
|
||||
|
||||
Returns:
|
||||
Wrapped handler function
|
||||
"""
|
||||
|
||||
@wraps(handler)
|
||||
async def wrapped(*args, **kwargs):
|
||||
# For project endpoints, always inject site filter
|
||||
# This locks all tools to the specific project's site
|
||||
if config.site_filter:
|
||||
# Extract site_id from project's site_filter (format: plugin_type_site_id)
|
||||
# The site parameter expects just the site identifier (site_id or alias)
|
||||
if "_" in config.site_filter:
|
||||
# site_filter is full_id like "wordpress_site1"
|
||||
parts = config.site_filter.split("_", 1)
|
||||
if len(parts) == 2:
|
||||
# Use the site_id part (site1)
|
||||
kwargs["site"] = parts[1]
|
||||
else:
|
||||
kwargs["site"] = config.site_filter
|
||||
else:
|
||||
kwargs["site"] = config.site_filter
|
||||
|
||||
# Call original handler
|
||||
return await handler(*args, **kwargs)
|
||||
|
||||
# Preserve function metadata
|
||||
wrapped.__name__ = tool_name
|
||||
wrapped.__doc__ = handler.__doc__
|
||||
|
||||
return wrapped
|
||||
|
||||
def get_endpoint(self, path: str) -> FastMCP | None:
|
||||
"""
|
||||
Get an endpoint by path.
|
||||
|
||||
Args:
|
||||
path: Endpoint path
|
||||
|
||||
Returns:
|
||||
FastMCP instance or None
|
||||
"""
|
||||
return self.endpoints.get(path)
|
||||
|
||||
def get_all_endpoints(self) -> dict[str, FastMCP]:
|
||||
"""Get all registered endpoints"""
|
||||
return self.endpoints.copy()
|
||||
|
||||
def get_endpoint_info(self) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Get information about all endpoints.
|
||||
|
||||
Returns:
|
||||
List of endpoint info dictionaries
|
||||
"""
|
||||
info = []
|
||||
for path, mcp in self.endpoints.items():
|
||||
# Get tool count
|
||||
# Note: This requires accessing FastMCP internals
|
||||
tool_count = len(mcp._tool_manager._tools) if hasattr(mcp, "_tool_manager") else 0
|
||||
|
||||
info.append(
|
||||
{
|
||||
"path": path,
|
||||
"name": mcp.name,
|
||||
"tool_count": tool_count,
|
||||
}
|
||||
)
|
||||
return info
|
||||
350
core/endpoints/middleware.py
Normal file
350
core/endpoints/middleware.py
Normal file
@@ -0,0 +1,350 @@
|
||||
"""
|
||||
Middleware for Multi-Endpoint Architecture
|
||||
|
||||
Provides authentication, rate limiting, and audit logging
|
||||
that works with the multi-endpoint architecture.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.server.dependencies import get_http_headers
|
||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||
|
||||
from core.api_keys import get_api_key_manager
|
||||
from core.audit_log import EventType, LogLevel, get_audit_logger
|
||||
from core.auth import get_auth_manager
|
||||
from core.context import clear_api_key_context, set_api_key_context
|
||||
from core.rate_limiter import get_rate_limiter
|
||||
|
||||
from .config import EndpointConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class AuthContext:
|
||||
"""Authentication context for a request"""
|
||||
|
||||
key_id: str | None = None
|
||||
project_id: str | None = None
|
||||
scope: str = "read"
|
||||
is_master_key: bool = False
|
||||
is_oauth_token: bool = False
|
||||
client_ip: str | None = None
|
||||
|
||||
class EndpointAuthMiddleware(Middleware):
|
||||
"""
|
||||
Authentication middleware for multi-endpoint architecture.
|
||||
|
||||
Validates API keys/tokens and enforces endpoint-specific access rules.
|
||||
"""
|
||||
|
||||
def __init__(self, endpoint_config: EndpointConfig):
|
||||
"""
|
||||
Initialize middleware with endpoint configuration.
|
||||
|
||||
Args:
|
||||
endpoint_config: Configuration for this endpoint
|
||||
"""
|
||||
self.config = endpoint_config
|
||||
self.auth_manager = get_auth_manager()
|
||||
self.api_key_manager = get_api_key_manager()
|
||||
|
||||
async def on_call_tool(self, context: MiddlewareContext, call_next: Callable):
|
||||
"""
|
||||
Handle tool call with authentication and authorization.
|
||||
|
||||
Args:
|
||||
context: Middleware context
|
||||
call_next: Next middleware in chain
|
||||
"""
|
||||
tool_name = getattr(context.message, "name", "unknown")
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Extract and validate authentication
|
||||
auth_context = await self._authenticate(context)
|
||||
|
||||
# Check endpoint access
|
||||
self._check_endpoint_access(auth_context)
|
||||
|
||||
# Check tool access
|
||||
self._check_tool_access(tool_name, auth_context)
|
||||
|
||||
# Set context for downstream handlers
|
||||
if auth_context.key_id:
|
||||
set_api_key_context(
|
||||
key_id=auth_context.key_id,
|
||||
project_id=auth_context.project_id or "*",
|
||||
scope=auth_context.scope,
|
||||
is_global=auth_context.project_id == "*",
|
||||
)
|
||||
|
||||
# Call the actual tool
|
||||
result = await call_next(context)
|
||||
|
||||
# Log success
|
||||
self._log_success(tool_name, auth_context, start_time)
|
||||
|
||||
return result
|
||||
|
||||
except ToolError:
|
||||
raise
|
||||
except Exception as e:
|
||||
self._log_error(tool_name, str(e), start_time)
|
||||
raise ToolError(f"Authentication error: {str(e)}")
|
||||
finally:
|
||||
clear_api_key_context()
|
||||
|
||||
async def _authenticate(self, context: MiddlewareContext) -> AuthContext:
|
||||
"""
|
||||
Extract and validate authentication from request.
|
||||
|
||||
Args:
|
||||
context: Middleware context
|
||||
|
||||
Returns:
|
||||
AuthContext with authentication details
|
||||
"""
|
||||
auth_context = AuthContext()
|
||||
|
||||
# Get headers
|
||||
try:
|
||||
headers = get_http_headers()
|
||||
except Exception:
|
||||
headers = {}
|
||||
|
||||
# Extract client IP
|
||||
auth_context.client_ip = headers.get("x-forwarded-for", "unknown")
|
||||
|
||||
# Get authorization header
|
||||
auth_header = headers.get("authorization", "")
|
||||
|
||||
if not auth_header:
|
||||
# No auth provided
|
||||
if self.config.require_master_key:
|
||||
raise ToolError("Master API key required for this endpoint")
|
||||
return auth_context
|
||||
|
||||
# Parse authorization
|
||||
if auth_header.startswith("Bearer "):
|
||||
token = auth_header[7:]
|
||||
else:
|
||||
token = auth_header
|
||||
|
||||
# Check token type
|
||||
if token.startswith("sk-"):
|
||||
# Master API key
|
||||
if self.auth_manager.validate_master_key(token):
|
||||
auth_context.is_master_key = True
|
||||
auth_context.project_id = "*"
|
||||
auth_context.scope = "admin"
|
||||
auth_context.key_id = "master"
|
||||
return auth_context
|
||||
else:
|
||||
raise ToolError("Invalid master API key")
|
||||
|
||||
elif token.startswith("cmp_"):
|
||||
# Project API key
|
||||
key = self.api_key_manager.get_key_by_token(token)
|
||||
if not key:
|
||||
raise ToolError("Invalid API key")
|
||||
|
||||
if key.revoked:
|
||||
raise ToolError("API key has been revoked")
|
||||
|
||||
if key.is_expired():
|
||||
raise ToolError("API key has expired")
|
||||
|
||||
auth_context.key_id = key.key_id
|
||||
auth_context.project_id = key.project_id
|
||||
auth_context.scope = key.scope
|
||||
return auth_context
|
||||
|
||||
else:
|
||||
# Possibly OAuth token (JWT)
|
||||
try:
|
||||
from core.oauth import get_token_manager
|
||||
|
||||
token_manager = get_token_manager()
|
||||
payload = token_manager.validate_access_token(token)
|
||||
|
||||
if payload:
|
||||
auth_context.is_oauth_token = True
|
||||
auth_context.project_id = payload.get("project_id", "*")
|
||||
auth_context.scope = payload.get("scope", "read")
|
||||
auth_context.key_id = f"oauth_{payload.get('sub', 'unknown')}"
|
||||
return auth_context
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise ToolError("Invalid authentication token")
|
||||
|
||||
def _check_endpoint_access(self, auth_context: AuthContext):
|
||||
"""
|
||||
Check if auth context allows access to this endpoint.
|
||||
|
||||
Args:
|
||||
auth_context: Authentication context
|
||||
"""
|
||||
# Master key always has access
|
||||
if auth_context.is_master_key:
|
||||
return
|
||||
|
||||
# Check if endpoint requires master key
|
||||
if self.config.require_master_key:
|
||||
raise ToolError(f"Endpoint {self.config.path} requires master API key")
|
||||
|
||||
# Check scope requirements
|
||||
if self.config.allowed_scopes:
|
||||
# Check if any of the user's scopes are allowed
|
||||
user_scopes = set(auth_context.scope.split())
|
||||
if not user_scopes & self.config.allowed_scopes:
|
||||
raise ToolError(
|
||||
f"Insufficient scope. Required: {self.config.allowed_scopes}, "
|
||||
f"Got: {user_scopes}"
|
||||
)
|
||||
|
||||
# Check plugin type access
|
||||
if auth_context.project_id and auth_context.project_id != "*":
|
||||
# Extract plugin type from project_id (e.g., "wordpress_site4" -> "wordpress")
|
||||
if "_" in auth_context.project_id:
|
||||
key_plugin_type = auth_context.project_id.split("_")[0]
|
||||
|
||||
# Check if endpoint allows this plugin type
|
||||
if self.config.plugin_types and key_plugin_type not in self.config.plugin_types:
|
||||
raise ToolError(
|
||||
f"API key for {key_plugin_type} cannot access "
|
||||
f"{self.config.endpoint_type.value} endpoint"
|
||||
)
|
||||
|
||||
def _check_tool_access(self, tool_name: str, auth_context: AuthContext):
|
||||
"""
|
||||
Check if auth context allows access to specific tool.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
auth_context: Authentication context
|
||||
"""
|
||||
# Master key has access to all tools
|
||||
if auth_context.is_master_key:
|
||||
return
|
||||
|
||||
# Check tool blacklist
|
||||
if not self.config.allows_tool(tool_name):
|
||||
raise ToolError(f"Access denied to tool: {tool_name}")
|
||||
|
||||
# Check site filter for project endpoints
|
||||
if self.config.site_filter:
|
||||
# Tool must be for the configured site
|
||||
# This is handled by parameter injection in the wrapper
|
||||
pass
|
||||
|
||||
def _log_success(self, tool_name: str, auth_context: AuthContext, start_time: float):
|
||||
"""Log successful tool execution"""
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
logger.debug(
|
||||
f"Tool {tool_name} executed successfully "
|
||||
f"(key={auth_context.key_id}, duration={duration_ms}ms)"
|
||||
)
|
||||
|
||||
def _log_error(self, tool_name: str, error: str, start_time: float):
|
||||
"""Log tool execution error"""
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
logger.warning(f"Tool {tool_name} failed: {error} (duration={duration_ms}ms)")
|
||||
|
||||
class EndpointRateLimitMiddleware(Middleware):
|
||||
"""
|
||||
Rate limiting middleware for multi-endpoint architecture.
|
||||
"""
|
||||
|
||||
def __init__(self, endpoint_config: EndpointConfig):
|
||||
self.config = endpoint_config
|
||||
self.rate_limiter = get_rate_limiter()
|
||||
|
||||
async def on_call_tool(self, context: MiddlewareContext, call_next: Callable):
|
||||
"""Apply rate limiting before tool execution"""
|
||||
# Get client identifier
|
||||
try:
|
||||
headers = get_http_headers()
|
||||
client_id = headers.get("authorization", "anonymous")[:50]
|
||||
except Exception:
|
||||
client_id = "unknown"
|
||||
|
||||
# Check rate limit
|
||||
allowed, info = self.rate_limiter.check_rate_limit(client_id)
|
||||
|
||||
if not allowed:
|
||||
raise ToolError(
|
||||
f"Rate limit exceeded. Retry after {info.get('retry_after', 60)} seconds"
|
||||
)
|
||||
|
||||
# Proceed with request
|
||||
return await call_next(context)
|
||||
|
||||
class EndpointAuditMiddleware(Middleware):
|
||||
"""
|
||||
Audit logging middleware for multi-endpoint architecture.
|
||||
"""
|
||||
|
||||
def __init__(self, endpoint_config: EndpointConfig):
|
||||
self.config = endpoint_config
|
||||
self.audit_logger = get_audit_logger()
|
||||
|
||||
async def on_call_tool(self, context: MiddlewareContext, call_next: Callable):
|
||||
"""Log tool execution to audit log"""
|
||||
tool_name = getattr(context.message, "name", "unknown")
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
result = await call_next(context)
|
||||
|
||||
# Log success
|
||||
self.audit_logger.log(
|
||||
level=LogLevel.INFO,
|
||||
event_type=EventType.TOOL_CALL,
|
||||
message=f"Tool executed: {tool_name}",
|
||||
details={
|
||||
"tool": tool_name,
|
||||
"endpoint": self.config.path,
|
||||
"duration_ms": int((time.time() - start_time) * 1000),
|
||||
"success": True,
|
||||
},
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# Log failure
|
||||
self.audit_logger.log(
|
||||
level=LogLevel.WARNING,
|
||||
event_type=EventType.TOOL_CALL,
|
||||
message=f"Tool failed: {tool_name}",
|
||||
details={
|
||||
"tool": tool_name,
|
||||
"endpoint": self.config.path,
|
||||
"duration_ms": int((time.time() - start_time) * 1000),
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
raise
|
||||
|
||||
def create_endpoint_middleware(endpoint_config: EndpointConfig) -> list:
|
||||
"""
|
||||
Create middleware stack for an endpoint.
|
||||
|
||||
Args:
|
||||
endpoint_config: Endpoint configuration
|
||||
|
||||
Returns:
|
||||
List of middleware instances
|
||||
"""
|
||||
return [
|
||||
EndpointAuthMiddleware(endpoint_config),
|
||||
EndpointRateLimitMiddleware(endpoint_config),
|
||||
EndpointAuditMiddleware(endpoint_config),
|
||||
]
|
||||
226
core/endpoints/registry.py
Normal file
226
core/endpoints/registry.py
Normal file
@@ -0,0 +1,226 @@
|
||||
"""
|
||||
Endpoint Registry
|
||||
|
||||
Central registry for managing all MCP endpoints.
|
||||
Handles routing requests to the correct endpoint based on path.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from starlette.routing import Mount, Route
|
||||
|
||||
from .config import (
|
||||
ENDPOINT_CONFIGS,
|
||||
EndpointConfig,
|
||||
EndpointType,
|
||||
create_project_endpoint_config,
|
||||
)
|
||||
from .factory import MCPEndpointFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class EndpointInfo:
|
||||
"""Information about a registered endpoint"""
|
||||
|
||||
path: str
|
||||
name: str
|
||||
description: str
|
||||
endpoint_type: EndpointType
|
||||
tool_count: int
|
||||
plugin_types: list[str]
|
||||
require_master_key: bool
|
||||
|
||||
class EndpointRegistry:
|
||||
"""
|
||||
Central registry for all MCP endpoints.
|
||||
|
||||
Manages endpoint creation, routing, and discovery.
|
||||
"""
|
||||
|
||||
def __init__(self, factory: MCPEndpointFactory):
|
||||
"""
|
||||
Initialize the endpoint registry.
|
||||
|
||||
Args:
|
||||
factory: Endpoint factory for creating MCP instances
|
||||
"""
|
||||
self.factory = factory
|
||||
self._endpoints: dict[str, FastMCP] = {}
|
||||
self._configs: dict[str, EndpointConfig] = {}
|
||||
self._initialized = False
|
||||
|
||||
def initialize_default_endpoints(self):
|
||||
"""
|
||||
Initialize the default set of endpoints.
|
||||
|
||||
Creates admin, wordpress, wordpress-advanced, and gitea endpoints.
|
||||
"""
|
||||
if self._initialized:
|
||||
logger.warning("Endpoints already initialized")
|
||||
return
|
||||
|
||||
logger.info("=" * 60)
|
||||
logger.info("Initializing Multi-Endpoint Architecture")
|
||||
logger.info("=" * 60)
|
||||
|
||||
# Create default endpoints
|
||||
for _endpoint_type, config in ENDPOINT_CONFIGS.items():
|
||||
try:
|
||||
mcp = self.factory.create_endpoint(config)
|
||||
self._endpoints[config.path] = mcp
|
||||
self._configs[config.path] = config
|
||||
logger.info(f" ✓ {config.path}: {config.name}")
|
||||
except Exception as e:
|
||||
logger.error(f" ✗ Failed to create {config.path}: {e}")
|
||||
|
||||
self._initialized = True
|
||||
|
||||
# Log summary
|
||||
self._log_summary()
|
||||
|
||||
def create_project_endpoint(
|
||||
self, project_id: str, plugin_type: str, site_alias: str | None = None
|
||||
) -> FastMCP:
|
||||
"""
|
||||
Create a dynamic endpoint for a specific project.
|
||||
|
||||
Args:
|
||||
project_id: Full project ID
|
||||
plugin_type: Plugin type
|
||||
site_alias: Optional site alias
|
||||
|
||||
Returns:
|
||||
Created FastMCP instance
|
||||
"""
|
||||
config = create_project_endpoint_config(project_id, plugin_type, site_alias)
|
||||
|
||||
# Check if already exists
|
||||
if config.path in self._endpoints:
|
||||
logger.info(f"Endpoint {config.path} already exists")
|
||||
return self._endpoints[config.path]
|
||||
|
||||
mcp = self.factory.create_endpoint(config)
|
||||
self._endpoints[config.path] = mcp
|
||||
self._configs[config.path] = config
|
||||
|
||||
logger.info(f"Created project endpoint: {config.path}")
|
||||
return mcp
|
||||
|
||||
def get_endpoint(self, path: str) -> FastMCP | None:
|
||||
"""
|
||||
Get endpoint by path.
|
||||
|
||||
Args:
|
||||
path: Endpoint path
|
||||
|
||||
Returns:
|
||||
FastMCP instance or None
|
||||
"""
|
||||
# Exact match
|
||||
if path in self._endpoints:
|
||||
return self._endpoints[path]
|
||||
|
||||
# Try with trailing slash
|
||||
if not path.endswith("/"):
|
||||
return self._endpoints.get(path + "/")
|
||||
|
||||
return None
|
||||
|
||||
def get_config(self, path: str) -> EndpointConfig | None:
|
||||
"""Get endpoint configuration by path"""
|
||||
return self._configs.get(path)
|
||||
|
||||
def list_endpoints(self) -> list[EndpointInfo]:
|
||||
"""
|
||||
List all registered endpoints with their info.
|
||||
|
||||
Returns:
|
||||
List of EndpointInfo objects
|
||||
"""
|
||||
endpoints = []
|
||||
|
||||
for path, config in self._configs.items():
|
||||
mcp = self._endpoints.get(path)
|
||||
tool_count = 0
|
||||
|
||||
if mcp:
|
||||
# Try to get tool count from FastMCP
|
||||
try:
|
||||
tool_count = len(mcp._tool_manager._tools)
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
endpoints.append(
|
||||
EndpointInfo(
|
||||
path=path,
|
||||
name=config.name,
|
||||
description=config.description,
|
||||
endpoint_type=config.endpoint_type,
|
||||
tool_count=tool_count,
|
||||
plugin_types=config.plugin_types,
|
||||
require_master_key=config.require_master_key,
|
||||
)
|
||||
)
|
||||
|
||||
return endpoints
|
||||
|
||||
def get_routes(self) -> list[Route]:
|
||||
"""
|
||||
Get Starlette routes for all endpoints.
|
||||
|
||||
Returns:
|
||||
List of Route objects for mounting
|
||||
"""
|
||||
routes = []
|
||||
|
||||
for path, mcp in self._endpoints.items():
|
||||
# Each MCP endpoint needs to handle its own routing
|
||||
# FastMCP provides an ASGI app
|
||||
routes.append(Mount(path, app=mcp.sse_app(), name=f"mcp_{path.replace('/', '_')}"))
|
||||
|
||||
return routes
|
||||
|
||||
def _log_summary(self):
|
||||
"""Log a summary of all endpoints"""
|
||||
logger.info("-" * 60)
|
||||
logger.info("Endpoint Summary:")
|
||||
logger.info("-" * 60)
|
||||
|
||||
total_tools = 0
|
||||
for info in self.list_endpoints():
|
||||
total_tools += info.tool_count
|
||||
auth_note = " (Master Key required)" if info.require_master_key else ""
|
||||
logger.info(f" {info.path}: {info.tool_count} tools{auth_note}")
|
||||
|
||||
logger.info("-" * 60)
|
||||
logger.info(f"Total: {len(self._endpoints)} endpoints")
|
||||
logger.info("=" * 60)
|
||||
|
||||
# Singleton instance
|
||||
_registry: EndpointRegistry | None = None
|
||||
|
||||
def get_endpoint_registry() -> EndpointRegistry:
|
||||
"""Get the global endpoint registry instance"""
|
||||
global _registry
|
||||
if _registry is None:
|
||||
raise RuntimeError(
|
||||
"Endpoint registry not initialized. " "Call initialize_endpoint_registry() first."
|
||||
)
|
||||
return _registry
|
||||
|
||||
def initialize_endpoint_registry(factory: MCPEndpointFactory) -> EndpointRegistry:
|
||||
"""
|
||||
Initialize the global endpoint registry.
|
||||
|
||||
Args:
|
||||
factory: Endpoint factory to use
|
||||
|
||||
Returns:
|
||||
Initialized EndpointRegistry
|
||||
"""
|
||||
global _registry
|
||||
_registry = EndpointRegistry(factory)
|
||||
return _registry
|
||||
683
core/health.py
Normal file
683
core/health.py
Normal file
@@ -0,0 +1,683 @@
|
||||
"""
|
||||
Enhanced Health Monitoring System for MCP Server (Phase 7.2)
|
||||
|
||||
This module provides comprehensive health monitoring capabilities including:
|
||||
- Response time tracking
|
||||
- Error rate monitoring
|
||||
- Historical metrics storage
|
||||
- Alert thresholds
|
||||
- Dependency health checks
|
||||
- System uptime tracking
|
||||
|
||||
Author: Coolify MCP Team
|
||||
Version: 7.2
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections import defaultdict, deque
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from core.audit_log import AuditLogger
|
||||
from core.project_manager import ProjectManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class HealthMetric:
|
||||
"""Individual health metric data point."""
|
||||
|
||||
timestamp: datetime
|
||||
project_id: str
|
||||
response_time_ms: float
|
||||
success: bool
|
||||
error_message: str | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return {
|
||||
"timestamp": self.timestamp.isoformat(),
|
||||
"project_id": self.project_id,
|
||||
"response_time_ms": self.response_time_ms,
|
||||
"success": self.success,
|
||||
"error_message": self.error_message,
|
||||
}
|
||||
|
||||
@dataclass
|
||||
class SystemMetrics:
|
||||
"""System-wide metrics."""
|
||||
|
||||
uptime_seconds: float
|
||||
total_requests: int
|
||||
successful_requests: int
|
||||
failed_requests: int
|
||||
average_response_time_ms: float
|
||||
error_rate_percent: float
|
||||
requests_per_minute: float
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dictionary."""
|
||||
return asdict(self)
|
||||
|
||||
@dataclass
|
||||
class ProjectHealthStatus:
|
||||
"""Comprehensive health status for a project."""
|
||||
|
||||
project_id: str
|
||||
healthy: bool
|
||||
last_check: datetime
|
||||
response_time_ms: float
|
||||
error_rate_percent: float
|
||||
recent_errors: list[str] = field(default_factory=list)
|
||||
alerts: list[str] = field(default_factory=list)
|
||||
details: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return {
|
||||
"project_id": self.project_id,
|
||||
"healthy": self.healthy,
|
||||
"last_check": self.last_check.isoformat(),
|
||||
"response_time_ms": self.response_time_ms,
|
||||
"error_rate_percent": self.error_rate_percent,
|
||||
"recent_errors": self.recent_errors,
|
||||
"alerts": self.alerts,
|
||||
"details": self.details,
|
||||
}
|
||||
|
||||
@dataclass
|
||||
class AlertThreshold:
|
||||
"""Alert threshold configuration."""
|
||||
|
||||
name: str
|
||||
metric: str # "response_time_ms", "error_rate_percent", etc.
|
||||
threshold: float
|
||||
comparison: str # "gt" (greater than), "lt" (less than), "eq" (equal)
|
||||
severity: str = "warning" # "info", "warning", "critical"
|
||||
|
||||
def check(self, value: float) -> bool:
|
||||
"""Check if value exceeds threshold."""
|
||||
if self.comparison == "gt":
|
||||
return value > self.threshold
|
||||
elif self.comparison == "lt":
|
||||
return value < self.threshold
|
||||
elif self.comparison == "eq":
|
||||
return value == self.threshold
|
||||
return False
|
||||
|
||||
class HealthMonitor:
|
||||
"""
|
||||
Enhanced health monitoring system with metrics tracking and alerting.
|
||||
|
||||
Features:
|
||||
- Real-time health checks
|
||||
- Response time tracking
|
||||
- Error rate monitoring
|
||||
- Historical metrics (last 24 hours)
|
||||
- Alert thresholds
|
||||
- System uptime tracking
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
project_manager: ProjectManager,
|
||||
audit_logger: AuditLogger | None = None,
|
||||
metrics_retention_hours: int = 24,
|
||||
max_metrics_per_project: int = 1000,
|
||||
):
|
||||
"""
|
||||
Initialize health monitor.
|
||||
|
||||
Args:
|
||||
project_manager: Project manager instance
|
||||
audit_logger: Optional audit logger for logging health events
|
||||
metrics_retention_hours: Hours to retain historical metrics
|
||||
max_metrics_per_project: Maximum metrics to store per project
|
||||
"""
|
||||
self.project_manager = project_manager
|
||||
self.audit_logger = audit_logger
|
||||
self.metrics_retention_hours = metrics_retention_hours
|
||||
self.max_metrics_per_project = max_metrics_per_project
|
||||
|
||||
# Metrics storage (in-memory)
|
||||
# Using deque for efficient FIFO operations
|
||||
self.metrics_history: dict[str, deque] = defaultdict(
|
||||
lambda: deque(maxlen=max_metrics_per_project)
|
||||
)
|
||||
|
||||
# Request counters
|
||||
self.total_requests = 0
|
||||
self.successful_requests = 0
|
||||
self.failed_requests = 0
|
||||
|
||||
# Response time tracking
|
||||
self.response_times: deque = deque(maxlen=1000) # Last 1000 requests
|
||||
|
||||
# System start time
|
||||
self.start_time = time.time()
|
||||
|
||||
# Alert thresholds (configurable)
|
||||
self.alert_thresholds: dict[str, list[AlertThreshold]] = defaultdict(list)
|
||||
self._setup_default_thresholds()
|
||||
|
||||
# Request rate tracking (for requests per minute)
|
||||
self.request_timestamps: deque = deque(maxlen=1000)
|
||||
|
||||
logger.info("HealthMonitor initialized (Phase 7.2)")
|
||||
|
||||
def _setup_default_thresholds(self):
|
||||
"""Setup default alert thresholds."""
|
||||
# Response time threshold: > 5000ms (5 seconds) is critical
|
||||
self.alert_thresholds["global"].append(
|
||||
AlertThreshold(
|
||||
name="High Response Time",
|
||||
metric="response_time_ms",
|
||||
threshold=5000.0,
|
||||
comparison="gt",
|
||||
severity="critical",
|
||||
)
|
||||
)
|
||||
|
||||
# Error rate threshold: > 10% is warning, > 25% is critical
|
||||
self.alert_thresholds["global"].append(
|
||||
AlertThreshold(
|
||||
name="High Error Rate",
|
||||
metric="error_rate_percent",
|
||||
threshold=10.0,
|
||||
comparison="gt",
|
||||
severity="warning",
|
||||
)
|
||||
)
|
||||
|
||||
self.alert_thresholds["global"].append(
|
||||
AlertThreshold(
|
||||
name="Critical Error Rate",
|
||||
metric="error_rate_percent",
|
||||
threshold=25.0,
|
||||
comparison="gt",
|
||||
severity="critical",
|
||||
)
|
||||
)
|
||||
|
||||
def add_alert_threshold(
|
||||
self,
|
||||
project_id: str,
|
||||
name: str,
|
||||
metric: str,
|
||||
threshold: float,
|
||||
comparison: str = "gt",
|
||||
severity: str = "warning",
|
||||
):
|
||||
"""
|
||||
Add a custom alert threshold for a project.
|
||||
|
||||
Args:
|
||||
project_id: Project ID or "global" for all projects
|
||||
name: Alert name
|
||||
metric: Metric to check
|
||||
threshold: Threshold value
|
||||
comparison: Comparison operator ("gt", "lt", "eq")
|
||||
severity: Alert severity ("info", "warning", "critical")
|
||||
"""
|
||||
alert = AlertThreshold(name, metric, threshold, comparison, severity)
|
||||
self.alert_thresholds[project_id].append(alert)
|
||||
logger.info(f"Added alert threshold '{name}' for {project_id}")
|
||||
|
||||
def record_request(
|
||||
self,
|
||||
project_id: str,
|
||||
response_time_ms: float,
|
||||
success: bool,
|
||||
error_message: str | None = None,
|
||||
):
|
||||
"""
|
||||
Record a request metric.
|
||||
|
||||
Args:
|
||||
project_id: Project that handled the request
|
||||
response_time_ms: Response time in milliseconds
|
||||
success: Whether request succeeded
|
||||
error_message: Error message if failed
|
||||
"""
|
||||
# Create metric
|
||||
metric = HealthMetric(
|
||||
timestamp=datetime.now(UTC),
|
||||
project_id=project_id,
|
||||
response_time_ms=response_time_ms,
|
||||
success=success,
|
||||
error_message=error_message,
|
||||
)
|
||||
|
||||
# Store in history
|
||||
self.metrics_history[project_id].append(metric)
|
||||
|
||||
# Update counters
|
||||
self.total_requests += 1
|
||||
if success:
|
||||
self.successful_requests += 1
|
||||
else:
|
||||
self.failed_requests += 1
|
||||
|
||||
# Track response time
|
||||
self.response_times.append(response_time_ms)
|
||||
|
||||
# Track request timestamp for rate calculation
|
||||
self.request_timestamps.append(time.time())
|
||||
|
||||
# Log to audit if available
|
||||
if self.audit_logger:
|
||||
self.audit_logger.log_system_event(
|
||||
event="health_metric_recorded",
|
||||
details={
|
||||
"project_id": project_id,
|
||||
"response_time_ms": response_time_ms,
|
||||
"success": success,
|
||||
"error_message": error_message,
|
||||
},
|
||||
)
|
||||
|
||||
def _cleanup_old_metrics(self, project_id: str):
|
||||
"""Remove metrics older than retention period."""
|
||||
if project_id not in self.metrics_history:
|
||||
return
|
||||
|
||||
cutoff_time = datetime.now(UTC) - timedelta(hours=self.metrics_retention_hours)
|
||||
metrics = self.metrics_history[project_id]
|
||||
|
||||
# Remove old metrics from the front of deque
|
||||
while metrics and metrics[0].timestamp < cutoff_time:
|
||||
metrics.popleft()
|
||||
|
||||
def get_project_metrics(self, project_id: str, hours: int = 1) -> dict[str, Any]:
|
||||
"""
|
||||
Get metrics for a specific project.
|
||||
|
||||
Args:
|
||||
project_id: Project ID
|
||||
hours: Number of hours of history to analyze
|
||||
|
||||
Returns:
|
||||
Dictionary with metrics
|
||||
"""
|
||||
self._cleanup_old_metrics(project_id)
|
||||
|
||||
if project_id not in self.metrics_history:
|
||||
return {"project_id": project_id, "total_requests": 0, "error": "No metrics available"}
|
||||
|
||||
# Filter metrics by time window
|
||||
cutoff_time = datetime.now(UTC) - timedelta(hours=hours)
|
||||
metrics = [m for m in self.metrics_history[project_id] if m.timestamp >= cutoff_time]
|
||||
|
||||
if not metrics:
|
||||
return {"project_id": project_id, "total_requests": 0, "time_window_hours": hours}
|
||||
|
||||
# Calculate statistics
|
||||
total_requests = len(metrics)
|
||||
successful = sum(1 for m in metrics if m.success)
|
||||
failed = total_requests - successful
|
||||
error_rate = (failed / total_requests * 100) if total_requests > 0 else 0.0
|
||||
|
||||
# Response time statistics
|
||||
response_times = [m.response_time_ms for m in metrics]
|
||||
avg_response = sum(response_times) / len(response_times) if response_times else 0.0
|
||||
min_response = min(response_times) if response_times else 0.0
|
||||
max_response = max(response_times) if response_times else 0.0
|
||||
|
||||
# Recent errors (last 5)
|
||||
recent_errors = [m.error_message for m in metrics if not m.success and m.error_message][-5:]
|
||||
|
||||
return {
|
||||
"project_id": project_id,
|
||||
"time_window_hours": hours,
|
||||
"total_requests": total_requests,
|
||||
"successful_requests": successful,
|
||||
"failed_requests": failed,
|
||||
"error_rate_percent": round(error_rate, 2),
|
||||
"response_time": {
|
||||
"average_ms": round(avg_response, 2),
|
||||
"min_ms": round(min_response, 2),
|
||||
"max_ms": round(max_response, 2),
|
||||
},
|
||||
"recent_errors": recent_errors,
|
||||
}
|
||||
|
||||
def _check_alerts(self, project_id: str, metrics: dict[str, Any]) -> list[str]:
|
||||
"""
|
||||
Check if any alert thresholds are exceeded.
|
||||
|
||||
Args:
|
||||
project_id: Project ID
|
||||
metrics: Current metrics
|
||||
|
||||
Returns:
|
||||
List of alert messages
|
||||
"""
|
||||
alerts = []
|
||||
|
||||
# Check global thresholds
|
||||
for threshold in self.alert_thresholds["global"]:
|
||||
if threshold.metric in metrics:
|
||||
value = metrics[threshold.metric]
|
||||
if threshold.check(value):
|
||||
alerts.append(
|
||||
f"[{threshold.severity.upper()}] {threshold.name}: "
|
||||
f"{threshold.metric}={value} (threshold: {threshold.threshold})"
|
||||
)
|
||||
|
||||
# Check project-specific thresholds
|
||||
for threshold in self.alert_thresholds.get(project_id, []):
|
||||
if threshold.metric in metrics:
|
||||
value = metrics[threshold.metric]
|
||||
if threshold.check(value):
|
||||
alerts.append(
|
||||
f"[{threshold.severity.upper()}] {threshold.name}: "
|
||||
f"{threshold.metric}={value} (threshold: {threshold.threshold})"
|
||||
)
|
||||
|
||||
return alerts
|
||||
|
||||
async def check_project_health(
|
||||
self, project_id: str, include_metrics: bool = True
|
||||
) -> ProjectHealthStatus:
|
||||
"""
|
||||
Perform comprehensive health check on a project.
|
||||
|
||||
Args:
|
||||
project_id: Project ID to check
|
||||
include_metrics: Whether to include historical metrics
|
||||
|
||||
Returns:
|
||||
ProjectHealthStatus object
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Get plugin instance
|
||||
plugin = self.project_manager.projects.get(project_id)
|
||||
if not plugin:
|
||||
return ProjectHealthStatus(
|
||||
project_id=project_id,
|
||||
healthy=False,
|
||||
last_check=datetime.now(UTC),
|
||||
response_time_ms=0.0,
|
||||
error_rate_percent=100.0,
|
||||
recent_errors=["Project not found"],
|
||||
alerts=["CRITICAL: Project not found"],
|
||||
)
|
||||
|
||||
# Perform health check
|
||||
health_result = await plugin.health_check()
|
||||
response_time_ms = (time.time() - start_time) * 1000
|
||||
|
||||
# Handle both dict and string (JSON) responses
|
||||
if isinstance(health_result, str):
|
||||
try:
|
||||
import json
|
||||
|
||||
health_result = json.loads(health_result)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# If not valid JSON, treat as error message
|
||||
health_result = {"healthy": False, "message": health_result}
|
||||
|
||||
# Ensure health_result is a dict
|
||||
if not isinstance(health_result, dict):
|
||||
health_result = {"healthy": False, "message": str(health_result)}
|
||||
|
||||
# Record this health check
|
||||
is_healthy = health_result.get("healthy", False) or health_result.get("success", False)
|
||||
self.record_request(
|
||||
project_id=project_id,
|
||||
response_time_ms=response_time_ms,
|
||||
success=is_healthy,
|
||||
error_message=(
|
||||
health_result.get("message") or health_result.get("error")
|
||||
if not is_healthy
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
# Get metrics if requested
|
||||
metrics_data = {}
|
||||
error_rate = 0.0
|
||||
recent_errors = []
|
||||
|
||||
if include_metrics:
|
||||
metrics_data = self.get_project_metrics(project_id, hours=1)
|
||||
error_rate = metrics_data.get("error_rate_percent", 0.0)
|
||||
recent_errors = metrics_data.get("recent_errors", [])
|
||||
|
||||
# Check alerts
|
||||
alert_check_data = {
|
||||
"response_time_ms": response_time_ms,
|
||||
"error_rate_percent": error_rate,
|
||||
}
|
||||
alerts = self._check_alerts(project_id, alert_check_data)
|
||||
|
||||
return ProjectHealthStatus(
|
||||
project_id=project_id,
|
||||
healthy=is_healthy,
|
||||
last_check=datetime.now(UTC),
|
||||
response_time_ms=response_time_ms,
|
||||
error_rate_percent=error_rate,
|
||||
recent_errors=recent_errors,
|
||||
alerts=alerts,
|
||||
details=health_result,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
response_time_ms = (time.time() - start_time) * 1000
|
||||
error_msg = str(e)
|
||||
|
||||
# Record failed health check
|
||||
self.record_request(
|
||||
project_id=project_id,
|
||||
response_time_ms=response_time_ms,
|
||||
success=False,
|
||||
error_message=error_msg,
|
||||
)
|
||||
|
||||
return ProjectHealthStatus(
|
||||
project_id=project_id,
|
||||
healthy=False,
|
||||
last_check=datetime.now(UTC),
|
||||
response_time_ms=response_time_ms,
|
||||
error_rate_percent=100.0,
|
||||
recent_errors=[error_msg],
|
||||
alerts=[f"CRITICAL: Health check failed - {error_msg}"],
|
||||
)
|
||||
|
||||
async def check_all_projects_health(self, include_metrics: bool = True) -> dict[str, Any]:
|
||||
"""
|
||||
Check health of all projects.
|
||||
|
||||
Args:
|
||||
include_metrics: Whether to include historical metrics
|
||||
|
||||
Returns:
|
||||
Dictionary with overall health status
|
||||
"""
|
||||
health_statuses = {}
|
||||
|
||||
# Check each project
|
||||
for project_id in self.project_manager.projects.keys():
|
||||
status = await self.check_project_health(project_id, include_metrics)
|
||||
health_statuses[project_id] = status.to_dict()
|
||||
|
||||
# Calculate summary
|
||||
total_projects = len(health_statuses)
|
||||
healthy_projects = sum(1 for s in health_statuses.values() if s["healthy"])
|
||||
unhealthy_projects = total_projects - healthy_projects
|
||||
|
||||
# Collect all alerts
|
||||
all_alerts = []
|
||||
for status in health_statuses.values():
|
||||
all_alerts.extend(status.get("alerts", []))
|
||||
|
||||
return {
|
||||
"timestamp": datetime.now(UTC).isoformat(),
|
||||
"status": (
|
||||
"healthy"
|
||||
if unhealthy_projects == 0
|
||||
else ("degraded" if healthy_projects > 0 else "unhealthy")
|
||||
),
|
||||
"summary": {
|
||||
"total_projects": total_projects,
|
||||
"healthy": healthy_projects,
|
||||
"unhealthy": unhealthy_projects,
|
||||
},
|
||||
"alerts": all_alerts,
|
||||
"projects": health_statuses,
|
||||
}
|
||||
|
||||
def get_system_metrics(self) -> SystemMetrics:
|
||||
"""
|
||||
Get overall system metrics.
|
||||
|
||||
Returns:
|
||||
SystemMetrics object
|
||||
"""
|
||||
# Calculate uptime
|
||||
uptime_seconds = time.time() - self.start_time
|
||||
|
||||
# Calculate average response time
|
||||
avg_response_time = (
|
||||
sum(self.response_times) / len(self.response_times) if self.response_times else 0.0
|
||||
)
|
||||
|
||||
# Calculate error rate
|
||||
error_rate = (
|
||||
(self.failed_requests / self.total_requests * 100) if self.total_requests > 0 else 0.0
|
||||
)
|
||||
|
||||
# Calculate requests per minute
|
||||
now = time.time()
|
||||
one_minute_ago = now - 60
|
||||
recent_requests = sum(1 for ts in self.request_timestamps if ts >= one_minute_ago)
|
||||
|
||||
return SystemMetrics(
|
||||
uptime_seconds=uptime_seconds,
|
||||
total_requests=self.total_requests,
|
||||
successful_requests=self.successful_requests,
|
||||
failed_requests=self.failed_requests,
|
||||
average_response_time_ms=round(avg_response_time, 2),
|
||||
error_rate_percent=round(error_rate, 2),
|
||||
requests_per_minute=recent_requests,
|
||||
)
|
||||
|
||||
def get_uptime(self) -> dict[str, Any]:
|
||||
"""
|
||||
Get system uptime information.
|
||||
|
||||
Returns:
|
||||
Dictionary with uptime details
|
||||
"""
|
||||
uptime_seconds = time.time() - self.start_time
|
||||
uptime_minutes = uptime_seconds / 60
|
||||
uptime_hours = uptime_minutes / 60
|
||||
uptime_days = uptime_hours / 24
|
||||
|
||||
return {
|
||||
"start_time": datetime.fromtimestamp(self.start_time, tz=UTC).isoformat(),
|
||||
"current_time": datetime.now(UTC).isoformat(),
|
||||
"uptime_seconds": round(uptime_seconds, 2),
|
||||
"uptime_minutes": round(uptime_minutes, 2),
|
||||
"uptime_hours": round(uptime_hours, 2),
|
||||
"uptime_days": round(uptime_days, 2),
|
||||
"uptime_formatted": self._format_uptime(uptime_seconds),
|
||||
}
|
||||
|
||||
def _format_uptime(self, seconds: float) -> str:
|
||||
"""Format uptime as human-readable string."""
|
||||
days = int(seconds // 86400)
|
||||
hours = int((seconds % 86400) // 3600)
|
||||
minutes = int((seconds % 3600) // 60)
|
||||
secs = int(seconds % 60)
|
||||
|
||||
parts = []
|
||||
if days > 0:
|
||||
parts.append(f"{days}d")
|
||||
if hours > 0:
|
||||
parts.append(f"{hours}h")
|
||||
if minutes > 0:
|
||||
parts.append(f"{minutes}m")
|
||||
parts.append(f"{secs}s")
|
||||
|
||||
return " ".join(parts)
|
||||
|
||||
def export_metrics(self, output_path: str | None = None, format: str = "json") -> str:
|
||||
"""
|
||||
Export all metrics to file.
|
||||
|
||||
Args:
|
||||
output_path: Output file path (default: logs/metrics_export.json)
|
||||
format: Export format ("json" only for now)
|
||||
|
||||
Returns:
|
||||
Path to exported file
|
||||
"""
|
||||
if output_path is None:
|
||||
output_path = "logs/metrics_export.json"
|
||||
|
||||
# Prepare export data
|
||||
export_data = {
|
||||
"export_time": datetime.now(UTC).isoformat(),
|
||||
"system_metrics": self.get_system_metrics().to_dict(),
|
||||
"uptime": self.get_uptime(),
|
||||
"projects": {},
|
||||
}
|
||||
|
||||
# Add per-project metrics
|
||||
for project_id in self.metrics_history.keys():
|
||||
export_data["projects"][project_id] = {
|
||||
"metrics": self.get_project_metrics(project_id, hours=24),
|
||||
"history": [m.to_dict() for m in self.metrics_history[project_id]],
|
||||
}
|
||||
|
||||
# Write to file
|
||||
output_file = Path(output_path)
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
json.dump(export_data, f, indent=2, ensure_ascii=False)
|
||||
|
||||
logger.info(f"Metrics exported to {output_path}")
|
||||
return str(output_file)
|
||||
|
||||
def reset_metrics(self):
|
||||
"""Reset all metrics (use with caution)."""
|
||||
self.metrics_history.clear()
|
||||
self.total_requests = 0
|
||||
self.successful_requests = 0
|
||||
self.failed_requests = 0
|
||||
self.response_times.clear()
|
||||
self.request_timestamps.clear()
|
||||
logger.warning("All metrics have been reset")
|
||||
|
||||
# Singleton instance
|
||||
_health_monitor: HealthMonitor | None = None
|
||||
|
||||
def get_health_monitor() -> HealthMonitor | None:
|
||||
"""Get the global health monitor instance."""
|
||||
return _health_monitor
|
||||
|
||||
def initialize_health_monitor(
|
||||
project_manager: ProjectManager, audit_logger: AuditLogger | None = None, **kwargs
|
||||
) -> HealthMonitor:
|
||||
"""
|
||||
Initialize the global health monitor.
|
||||
|
||||
Args:
|
||||
project_manager: Project manager instance
|
||||
audit_logger: Optional audit logger
|
||||
**kwargs: Additional configuration options
|
||||
|
||||
Returns:
|
||||
HealthMonitor instance
|
||||
"""
|
||||
global _health_monitor
|
||||
_health_monitor = HealthMonitor(project_manager, audit_logger, **kwargs)
|
||||
return _health_monitor
|
||||
194
core/i18n.py
Normal file
194
core/i18n.py
Normal file
@@ -0,0 +1,194 @@
|
||||
"""
|
||||
Internationalization (i18n) utilities for OAuth Authorization pages.
|
||||
|
||||
Supports English (en) and Persian/Farsi (fa) languages.
|
||||
Language is auto-detected from Accept-Language header or query parameter.
|
||||
"""
|
||||
|
||||
# Translation dictionary for OAuth authorization pages
|
||||
TRANSLATIONS: dict[str, dict[str, str]] = {
|
||||
"en": {
|
||||
# Page title
|
||||
"page_title": "Authorization Required - MCP Hub",
|
||||
# Header section
|
||||
"auth_required": "Authorization Required",
|
||||
"wants_access": "{client_name} wants to access your MCP Hub",
|
||||
# Client information
|
||||
"client_info": "Client Information",
|
||||
"client_id_label": "Client ID:",
|
||||
"redirect_uri_label": "Redirect URI:",
|
||||
# Permissions section
|
||||
"requested_permissions": "Requested Permissions:",
|
||||
"no_permissions": "No specific permissions requested",
|
||||
# Form section
|
||||
"enter_api_key": "Enter your API Key to authorize",
|
||||
"api_key_placeholder": "cmp_xxxxxxxxxxxxxxxx",
|
||||
"api_key_note": "Your API key will be validated and used to grant permissions to this application.",
|
||||
# Buttons
|
||||
"approve": "Authorize",
|
||||
"deny": "Deny",
|
||||
# Security notice
|
||||
"security_note": "Security Note:",
|
||||
"security_message": "Only authorize applications you trust. The OAuth token will inherit the permissions from your API key.",
|
||||
# Error page
|
||||
"error_title": "Authorization Error - MCP Hub",
|
||||
"auth_error": "Authorization Error",
|
||||
"unable_to_complete": "Unable to complete the authorization request",
|
||||
"error_code": "Error Code:",
|
||||
"error_description": "Description:",
|
||||
"common_solutions": "Common Solutions:",
|
||||
"solution_1": "Verify your API key is correct and active",
|
||||
"solution_2": "Check that the client application is properly configured",
|
||||
"solution_3": "Ensure the redirect URI matches the registered URI",
|
||||
"return_to_app": "Return to Application",
|
||||
"close_window": "Close Window",
|
||||
"need_help": "Need help? Check the",
|
||||
"documentation": "documentation",
|
||||
# Footer
|
||||
"secured_with": "This connection is secured with OAuth 2.1 + PKCE",
|
||||
},
|
||||
"fa": {
|
||||
# Page title
|
||||
"page_title": "احراز هویت مورد نیاز - MCP Hub",
|
||||
# Header section
|
||||
"auth_required": "احراز هویت مورد نیاز",
|
||||
"wants_access": "{client_name} میخواهد به MCP Hub شما دسترسی داشته باشد",
|
||||
# Client information
|
||||
"client_info": "اطلاعات برنامه",
|
||||
"client_id_label": "شناسه برنامه:",
|
||||
"redirect_uri_label": "آدرس بازگشت:",
|
||||
# Permissions section
|
||||
"requested_permissions": "دسترسیهای درخواستی:",
|
||||
"no_permissions": "دسترسی خاصی درخواست نشده",
|
||||
# Form section
|
||||
"enter_api_key": "API Key خود را برای احراز هویت وارد کنید",
|
||||
"api_key_placeholder": "cmp_xxxxxxxxxxxxxxxx",
|
||||
"api_key_note": "API Key شما اعتبارسنجی شده و برای اعطای دسترسی به این برنامه استفاده خواهد شد.",
|
||||
# Buttons
|
||||
"approve": "تایید",
|
||||
"deny": "رد",
|
||||
# Security notice
|
||||
"security_note": "نکته امنیتی:",
|
||||
"security_message": "فقط برنامههایی را که به آنها اعتماد دارید تایید کنید. توکن OAuth دسترسیهای API Key شما را به ارث خواهد برد.",
|
||||
# Error page
|
||||
"error_title": "خطای احراز هویت - MCP Hub",
|
||||
"auth_error": "خطای احراز هویت",
|
||||
"unable_to_complete": "امکان تکمیل درخواست احراز هویت وجود ندارد",
|
||||
"error_code": "کد خطا:",
|
||||
"error_description": "توضیحات:",
|
||||
"common_solutions": "راهحلهای رایج:",
|
||||
"solution_1": "اطمینان حاصل کنید که API Key شما صحیح و فعال است",
|
||||
"solution_2": "بررسی کنید که برنامه به درستی پیکربندی شده است",
|
||||
"solution_3": "اطمینان حاصل کنید که آدرس بازگشت با آدرس ثبتشده مطابقت دارد",
|
||||
"return_to_app": "بازگشت به برنامه",
|
||||
"close_window": "بستن پنجره",
|
||||
"need_help": "نیاز به کمک دارید؟",
|
||||
"documentation": "مستندات",
|
||||
# Footer
|
||||
"secured_with": "این اتصال با OAuth 2.1 + PKCE امن شده است",
|
||||
},
|
||||
}
|
||||
|
||||
def detect_language(accept_language: str | None = None, query_lang: str | None = None) -> str:
|
||||
"""
|
||||
Detect user's preferred language from Accept-Language header or query parameter.
|
||||
|
||||
Args:
|
||||
accept_language: Accept-Language header value (e.g., "en-US,en;q=0.9,fa;q=0.8")
|
||||
query_lang: Explicit language parameter from query string (takes priority)
|
||||
|
||||
Returns:
|
||||
Language code: "en" or "fa"
|
||||
"""
|
||||
# Priority 1: Explicit query parameter
|
||||
if query_lang:
|
||||
lang = query_lang.lower().strip()
|
||||
if lang in ["fa", "persian", "farsi"]:
|
||||
return "fa"
|
||||
if lang in ["en", "english"]:
|
||||
return "en"
|
||||
|
||||
# Priority 2: Accept-Language header
|
||||
if accept_language:
|
||||
# Parse Accept-Language header
|
||||
# Format: "en-US,en;q=0.9,fa;q=0.8,fa-IR;q=0.7"
|
||||
languages = []
|
||||
for lang_entry in accept_language.split(","):
|
||||
# Extract language code (ignore quality value)
|
||||
lang = lang_entry.split(";")[0].strip().lower()
|
||||
languages.append(lang)
|
||||
|
||||
# Check for Persian/Farsi
|
||||
for lang in languages:
|
||||
if lang.startswith("fa"): # fa, fa-IR, fa-AF
|
||||
return "fa"
|
||||
|
||||
# Check for English
|
||||
for lang in languages:
|
||||
if lang.startswith("en"): # en, en-US, en-GB
|
||||
return "en"
|
||||
|
||||
# Default: English
|
||||
return "en"
|
||||
|
||||
def get_translation(lang: str, key: str, **kwargs) -> str:
|
||||
"""
|
||||
Get translated string for the given language and key.
|
||||
|
||||
Args:
|
||||
lang: Language code ("en" or "fa")
|
||||
key: Translation key
|
||||
**kwargs: Format arguments for string interpolation
|
||||
|
||||
Returns:
|
||||
Translated string with format arguments applied
|
||||
|
||||
Example:
|
||||
>>> get_translation("en", "wants_access", client_name="ChatGPT")
|
||||
'ChatGPT wants to access your MCP Hub'
|
||||
|
||||
>>> get_translation("fa", "wants_access", client_name="ChatGPT")
|
||||
'ChatGPT میخواهد به MCP Hub شما دسترسی داشته باشد'
|
||||
"""
|
||||
# Get language dictionary (fallback to English)
|
||||
lang_dict = TRANSLATIONS.get(lang, TRANSLATIONS["en"])
|
||||
|
||||
# Get translation (fallback to key itself)
|
||||
text = lang_dict.get(key, key)
|
||||
|
||||
# Apply format arguments if provided
|
||||
if kwargs:
|
||||
try:
|
||||
return text.format(**kwargs)
|
||||
except (KeyError, ValueError):
|
||||
# If formatting fails, return raw text
|
||||
return text
|
||||
|
||||
return text
|
||||
|
||||
def get_all_translations(lang: str) -> dict[str, str]:
|
||||
"""
|
||||
Get all translations for a language as a dictionary.
|
||||
|
||||
Useful for passing to templates.
|
||||
|
||||
Args:
|
||||
lang: Language code ("en" or "fa")
|
||||
|
||||
Returns:
|
||||
Dictionary of all translations for the language
|
||||
"""
|
||||
return TRANSLATIONS.get(lang, TRANSLATIONS["en"])
|
||||
|
||||
def get_language_name(lang: str) -> str:
|
||||
"""
|
||||
Get human-readable language name.
|
||||
|
||||
Args:
|
||||
lang: Language code
|
||||
|
||||
Returns:
|
||||
Language name in its native form
|
||||
"""
|
||||
names = {"en": "English", "fa": "فارسی"}
|
||||
return names.get(lang, "English")
|
||||
12
core/middleware/__init__.py
Normal file
12
core/middleware/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
"""
|
||||
Middleware Package
|
||||
|
||||
Organized middleware for MCP server.
|
||||
Part of Option B clean architecture refactoring.
|
||||
"""
|
||||
|
||||
from core.middleware.audit import AuditLoggingMiddleware
|
||||
from core.middleware.auth import UserAuthMiddleware
|
||||
from core.middleware.rate_limit import RateLimitMiddleware
|
||||
|
||||
__all__ = ["UserAuthMiddleware", "AuditLoggingMiddleware", "RateLimitMiddleware"]
|
||||
65
core/oauth/__init__.py
Normal file
65
core/oauth/__init__.py
Normal file
@@ -0,0 +1,65 @@
|
||||
"""
|
||||
OAuth 2.1 Infrastructure for MCP Hub
|
||||
|
||||
This module provides OAuth 2.1 compliant authentication and authorization
|
||||
infrastructure with PKCE support, refresh token rotation, and security best practices.
|
||||
"""
|
||||
|
||||
# Schemas
|
||||
# Client Registry
|
||||
from .client_registry import ClientRegistry, get_client_registry
|
||||
|
||||
# CSRF Protection (Phase E)
|
||||
from .csrf import CSRFTokenManager, get_csrf_manager
|
||||
|
||||
# PKCE utilities
|
||||
from .pkce import generate_code_challenge, generate_code_verifier, validate_code_challenge
|
||||
from .schemas import (
|
||||
AccessToken,
|
||||
AuthorizationCode,
|
||||
OAuthClient,
|
||||
RefreshToken,
|
||||
TokenRequest,
|
||||
TokenResponse,
|
||||
)
|
||||
|
||||
# OAuth Server
|
||||
from .server import OAuthError, OAuthServer, get_oauth_server
|
||||
|
||||
# Storage
|
||||
from .storage import BaseStorage, JSONStorage, get_storage
|
||||
|
||||
# Token Manager
|
||||
from .token_manager import SecurityError, TokenManager, get_token_manager
|
||||
|
||||
__all__ = [
|
||||
# Schemas
|
||||
"OAuthClient",
|
||||
"AuthorizationCode",
|
||||
"AccessToken",
|
||||
"RefreshToken",
|
||||
"TokenRequest",
|
||||
"TokenResponse",
|
||||
# PKCE
|
||||
"generate_code_verifier",
|
||||
"generate_code_challenge",
|
||||
"validate_code_challenge",
|
||||
# Storage
|
||||
"BaseStorage",
|
||||
"JSONStorage",
|
||||
"get_storage",
|
||||
# Client Registry
|
||||
"ClientRegistry",
|
||||
"get_client_registry",
|
||||
# Token Manager
|
||||
"TokenManager",
|
||||
"SecurityError",
|
||||
"get_token_manager",
|
||||
# OAuth Server
|
||||
"OAuthServer",
|
||||
"OAuthError",
|
||||
"get_oauth_server",
|
||||
# CSRF Protection
|
||||
"CSRFTokenManager",
|
||||
"get_csrf_manager",
|
||||
]
|
||||
154
core/oauth/client_registry.py
Normal file
154
core/oauth/client_registry.py
Normal file
@@ -0,0 +1,154 @@
|
||||
"""
|
||||
OAuth 2.1 Client Registry
|
||||
Manages OAuth client registration and validation
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import secrets
|
||||
from pathlib import Path
|
||||
|
||||
from .schemas import OAuthClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ClientRegistry:
|
||||
"""
|
||||
OAuth Client Registry with JSON storage.
|
||||
|
||||
Storage: data/oauth_clients.json
|
||||
"""
|
||||
|
||||
def __init__(self, data_dir: str = "/app/data"):
|
||||
self.data_dir = Path(data_dir)
|
||||
self.data_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.clients_file = self.data_dir / "oauth_clients.json"
|
||||
|
||||
# Initialize if not exists
|
||||
if not self.clients_file.exists():
|
||||
self._write_clients({})
|
||||
|
||||
# Load clients into memory (small dataset)
|
||||
self.clients: dict[str, OAuthClient] = self._load_clients()
|
||||
|
||||
def _read_clients(self) -> dict:
|
||||
"""Read clients JSON file"""
|
||||
try:
|
||||
with open(self.clients_file) as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
logger.error(f"Error reading clients file: {e}")
|
||||
return {}
|
||||
|
||||
def _write_clients(self, data: dict) -> bool:
|
||||
"""Write clients JSON file"""
|
||||
try:
|
||||
with open(self.clients_file, "w") as f:
|
||||
json.dump(data, f, indent=2, default=str)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error writing clients file: {e}")
|
||||
return False
|
||||
|
||||
def _load_clients(self) -> dict[str, OAuthClient]:
|
||||
"""Load clients from file"""
|
||||
data = self._read_clients()
|
||||
return {client_id: OAuthClient(**client_data) for client_id, client_data in data.items()}
|
||||
|
||||
def _save_clients(self) -> bool:
|
||||
"""Save clients to file"""
|
||||
data = {
|
||||
client_id: client.model_dump(mode="json") for client_id, client in self.clients.items()
|
||||
}
|
||||
return self._write_clients(data)
|
||||
|
||||
def create_client(
|
||||
self,
|
||||
client_name: str,
|
||||
redirect_uris: list[str],
|
||||
grant_types: list[str] | None = None,
|
||||
allowed_scopes: list[str] | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> tuple[str, str]:
|
||||
"""
|
||||
Create new OAuth client.
|
||||
|
||||
Returns:
|
||||
(client_id, client_secret) tuple
|
||||
"""
|
||||
# Generate client_id and client_secret
|
||||
client_id = f"cmp_client_{secrets.token_urlsafe(16)}"
|
||||
client_secret = secrets.token_urlsafe(32)
|
||||
|
||||
# Hash client secret
|
||||
client_secret_hash = hashlib.sha256(client_secret.encode()).hexdigest()
|
||||
|
||||
# Create client
|
||||
client = OAuthClient(
|
||||
client_id=client_id,
|
||||
client_secret_hash=client_secret_hash,
|
||||
client_name=client_name,
|
||||
redirect_uris=redirect_uris,
|
||||
grant_types=grant_types or ["authorization_code", "refresh_token"],
|
||||
allowed_scopes=allowed_scopes or ["read", "write"],
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
# Save
|
||||
self.clients[client_id] = client
|
||||
self._save_clients()
|
||||
|
||||
logger.info(f"Created OAuth client: {client_id} ({client_name})")
|
||||
|
||||
# Return client_secret in plain text (only time it's visible)
|
||||
return client_id, client_secret
|
||||
|
||||
def get_client(self, client_id: str) -> OAuthClient | None:
|
||||
"""Get client by ID"""
|
||||
return self.clients.get(client_id)
|
||||
|
||||
def list_clients(self) -> list[OAuthClient]:
|
||||
"""List all clients"""
|
||||
return list(self.clients.values())
|
||||
|
||||
def validate_client_secret(self, client_id: str, client_secret: str) -> bool:
|
||||
"""
|
||||
Validate client secret.
|
||||
|
||||
Args:
|
||||
client_id: Client ID
|
||||
client_secret: Plain text client secret
|
||||
|
||||
Returns:
|
||||
True if valid, False otherwise
|
||||
"""
|
||||
client = self.get_client(client_id)
|
||||
if not client:
|
||||
return False
|
||||
|
||||
# Hash provided secret
|
||||
secret_hash = hashlib.sha256(client_secret.encode()).hexdigest()
|
||||
|
||||
# Constant-time comparison
|
||||
return secrets.compare_digest(secret_hash, client.client_secret_hash)
|
||||
|
||||
def delete_client(self, client_id: str) -> bool:
|
||||
"""Delete OAuth client"""
|
||||
if client_id in self.clients:
|
||||
del self.clients[client_id]
|
||||
self._save_clients()
|
||||
logger.info(f"Deleted OAuth client: {client_id}")
|
||||
return True
|
||||
return False
|
||||
|
||||
# Singleton instance
|
||||
_client_registry: ClientRegistry | None = None
|
||||
|
||||
def get_client_registry() -> ClientRegistry:
|
||||
"""Get singleton ClientRegistry instance"""
|
||||
global _client_registry
|
||||
if _client_registry is None:
|
||||
_client_registry = ClientRegistry()
|
||||
return _client_registry
|
||||
123
core/oauth/csrf.py
Normal file
123
core/oauth/csrf.py
Normal file
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
CSRF Protection for OAuth Authorization
|
||||
|
||||
Implements token-based CSRF protection for the OAuth authorization flow.
|
||||
Tokens are stored in-memory with 10-minute expiry.
|
||||
|
||||
Part of Phase E: Custom OAuth Authorization Page
|
||||
"""
|
||||
|
||||
import secrets
|
||||
import time
|
||||
|
||||
class CSRFTokenManager:
|
||||
"""
|
||||
Manages CSRF tokens for OAuth authorization requests.
|
||||
|
||||
Features:
|
||||
- Generates cryptographically secure random tokens
|
||||
- In-memory storage with automatic expiry
|
||||
- 10-minute token lifetime
|
||||
- Automatic cleanup of expired tokens
|
||||
"""
|
||||
|
||||
def __init__(self, token_lifetime_seconds: int = 600):
|
||||
"""
|
||||
Initialize CSRF token manager.
|
||||
|
||||
Args:
|
||||
token_lifetime_seconds: Token lifetime in seconds (default: 600 = 10 minutes)
|
||||
"""
|
||||
self._tokens: dict[str, float] = {} # token -> expiry_timestamp
|
||||
self._token_lifetime = token_lifetime_seconds
|
||||
|
||||
def generate_token(self) -> str:
|
||||
"""
|
||||
Generate a new CSRF token.
|
||||
|
||||
Returns:
|
||||
Secure random token (32 bytes, hex-encoded = 64 characters)
|
||||
"""
|
||||
token = secrets.token_hex(32)
|
||||
expiry = time.time() + self._token_lifetime
|
||||
self._tokens[token] = expiry
|
||||
|
||||
# Cleanup expired tokens (lazy cleanup)
|
||||
self._cleanup_expired()
|
||||
|
||||
return token
|
||||
|
||||
def validate_token(self, token: str, consume: bool = True) -> bool:
|
||||
"""
|
||||
Validate a CSRF token.
|
||||
|
||||
Args:
|
||||
token: CSRF token to validate
|
||||
consume: If True, token is removed after validation (one-time use)
|
||||
|
||||
Returns:
|
||||
True if token is valid and not expired, False otherwise
|
||||
"""
|
||||
if not token or token not in self._tokens:
|
||||
return False
|
||||
|
||||
expiry = self._tokens[token]
|
||||
now = time.time()
|
||||
|
||||
# Check if token is expired
|
||||
if now > expiry:
|
||||
# Remove expired token
|
||||
self._tokens.pop(token, None)
|
||||
return False
|
||||
|
||||
# Token is valid
|
||||
if consume:
|
||||
# Remove token (one-time use)
|
||||
self._tokens.pop(token, None)
|
||||
|
||||
return True
|
||||
|
||||
def _cleanup_expired(self):
|
||||
"""
|
||||
Remove expired tokens from storage.
|
||||
|
||||
Called automatically during token generation to prevent memory leaks.
|
||||
"""
|
||||
now = time.time()
|
||||
expired_tokens = [token for token, expiry in self._tokens.items() if now > expiry]
|
||||
|
||||
for token in expired_tokens:
|
||||
self._tokens.pop(token, None)
|
||||
|
||||
def get_stats(self) -> dict[str, int]:
|
||||
"""
|
||||
Get statistics about stored tokens.
|
||||
|
||||
Returns:
|
||||
Dictionary with token statistics
|
||||
"""
|
||||
now = time.time()
|
||||
active_tokens = sum(1 for expiry in self._tokens.values() if expiry > now)
|
||||
expired_tokens = len(self._tokens) - active_tokens
|
||||
|
||||
return {
|
||||
"total_tokens": len(self._tokens),
|
||||
"active_tokens": active_tokens,
|
||||
"expired_tokens": expired_tokens,
|
||||
"token_lifetime_seconds": self._token_lifetime,
|
||||
}
|
||||
|
||||
# Global CSRF token manager instance
|
||||
_csrf_manager: CSRFTokenManager | None = None
|
||||
|
||||
def get_csrf_manager() -> CSRFTokenManager:
|
||||
"""
|
||||
Get the global CSRF token manager instance.
|
||||
|
||||
Returns:
|
||||
Global CSRFTokenManager instance (singleton)
|
||||
"""
|
||||
global _csrf_manager
|
||||
if _csrf_manager is None:
|
||||
_csrf_manager = CSRFTokenManager()
|
||||
return _csrf_manager
|
||||
91
core/oauth/pkce.py
Normal file
91
core/oauth/pkce.py
Normal file
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
PKCE (Proof Key for Code Exchange) Implementation
|
||||
RFC 7636 - OAuth 2.1 Mandatory
|
||||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import secrets
|
||||
from typing import Literal
|
||||
|
||||
def generate_code_verifier(length: int = 64) -> str:
|
||||
"""
|
||||
Generate PKCE code verifier.
|
||||
|
||||
Args:
|
||||
length: Length of verifier (43-128 characters)
|
||||
|
||||
Returns:
|
||||
URL-safe random string
|
||||
|
||||
OAuth 2.1 Spec:
|
||||
- Length: 43-128 characters
|
||||
- Character set: [A-Za-z0-9-._~] (unreserved characters)
|
||||
"""
|
||||
if not 43 <= length <= 128:
|
||||
raise ValueError("Code verifier length must be between 43-128")
|
||||
|
||||
# Generate random bytes and encode as URL-safe base64
|
||||
random_bytes = secrets.token_bytes(length)
|
||||
verifier = base64.urlsafe_b64encode(random_bytes).decode("utf-8")
|
||||
|
||||
# Remove padding and truncate to desired length
|
||||
verifier = verifier.rstrip("=")[:length]
|
||||
|
||||
return verifier
|
||||
|
||||
def generate_code_challenge(code_verifier: str, method: Literal["S256"] = "S256") -> str:
|
||||
"""
|
||||
Generate PKCE code challenge from verifier.
|
||||
|
||||
Args:
|
||||
code_verifier: Code verifier string
|
||||
method: Challenge method (OAuth 2.1: only S256)
|
||||
|
||||
Returns:
|
||||
Base64 URL-encoded SHA256 hash of verifier
|
||||
|
||||
OAuth 2.1 Spec:
|
||||
- Only S256 method is allowed (plain is removed)
|
||||
- code_challenge = BASE64URL(SHA256(code_verifier))
|
||||
"""
|
||||
if method != "S256":
|
||||
raise ValueError("OAuth 2.1 only supports S256 method (plain is removed)")
|
||||
|
||||
if not code_verifier:
|
||||
raise ValueError("Code verifier cannot be empty")
|
||||
|
||||
# SHA256 hash
|
||||
digest = hashlib.sha256(code_verifier.encode("utf-8")).digest()
|
||||
|
||||
# Base64 URL encode (no padding)
|
||||
challenge = base64.urlsafe_b64encode(digest).decode("utf-8").rstrip("=")
|
||||
|
||||
return challenge
|
||||
|
||||
def validate_code_challenge(
|
||||
code_verifier: str, code_challenge: str, method: Literal["S256"] = "S256"
|
||||
) -> bool:
|
||||
"""
|
||||
Validate PKCE code verifier against challenge.
|
||||
|
||||
Args:
|
||||
code_verifier: Code verifier from client
|
||||
code_challenge: Code challenge from authorization request
|
||||
method: Challenge method
|
||||
|
||||
Returns:
|
||||
True if valid, False otherwise
|
||||
"""
|
||||
if method != "S256":
|
||||
raise ValueError("Only S256 method is supported")
|
||||
|
||||
try:
|
||||
# Generate challenge from verifier
|
||||
expected_challenge = generate_code_challenge(code_verifier, method)
|
||||
|
||||
# Constant-time comparison to prevent timing attacks
|
||||
return secrets.compare_digest(expected_challenge, code_challenge)
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
145
core/oauth/schemas.py
Normal file
145
core/oauth/schemas.py
Normal file
@@ -0,0 +1,145 @@
|
||||
"""
|
||||
OAuth 2.1 Pydantic Schemas
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
class OAuthClient(BaseModel):
|
||||
"""OAuth 2.1 Client Model"""
|
||||
|
||||
client_id: str = Field(..., description="Client identifier (e.g., cmp_client_xxx)")
|
||||
client_secret_hash: str = Field(..., description="SHA256 hash of client secret")
|
||||
client_name: str = Field(..., description="Human-readable client name")
|
||||
redirect_uris: list[str] = Field(..., description="Allowed redirect URIs (exact match)")
|
||||
grant_types: list[str] = Field(
|
||||
default=["authorization_code", "refresh_token"], description="Allowed grant types"
|
||||
)
|
||||
response_types: list[str] = Field(default=["code"], description="Allowed response types")
|
||||
scope: str = Field(default="read", description="Default scope for this client")
|
||||
allowed_scopes: list[str] = Field(
|
||||
default=["read", "write"], description="All scopes this client can request"
|
||||
)
|
||||
token_endpoint_auth_method: str = Field(
|
||||
default="client_secret_post", description="Token endpoint authentication method"
|
||||
)
|
||||
created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||
metadata: dict = Field(default_factory=dict)
|
||||
|
||||
@field_validator("redirect_uris")
|
||||
def validate_redirect_uris(cls, v):
|
||||
"""Validate redirect URIs format"""
|
||||
for uri in v:
|
||||
if not uri.startswith(("http://", "https://")):
|
||||
raise ValueError(f"Invalid redirect URI: {uri}")
|
||||
return v
|
||||
|
||||
@field_validator("grant_types")
|
||||
def validate_grant_types(cls, v):
|
||||
"""Validate grant types"""
|
||||
allowed = ["authorization_code", "refresh_token", "client_credentials"]
|
||||
for grant in v:
|
||||
if grant not in allowed:
|
||||
raise ValueError(f"Invalid grant type: {grant}")
|
||||
return v
|
||||
|
||||
class AuthorizationCode(BaseModel):
|
||||
"""Authorization Code Model"""
|
||||
|
||||
code: str = Field(..., description="Authorization code (token_urlsafe)")
|
||||
client_id: str
|
||||
redirect_uri: str
|
||||
scope: str
|
||||
code_challenge: str = Field(..., description="PKCE code challenge")
|
||||
code_challenge_method: str = Field(default="S256")
|
||||
expires_at: datetime
|
||||
used: bool = Field(default=False)
|
||||
user_id: str | None = None # If user-based auth
|
||||
|
||||
# API Key-based authorization
|
||||
api_key_id: str | None = None # API Key ID for scope/project inheritance
|
||||
api_key_project_id: str | None = None # Project ID from API Key
|
||||
api_key_scope: str | None = None # Scope from API Key
|
||||
|
||||
def is_expired(self) -> bool:
|
||||
"""Check if code is expired"""
|
||||
now = datetime.now(UTC)
|
||||
expires = self.expires_at
|
||||
if expires.tzinfo is None:
|
||||
expires = expires.replace(tzinfo=UTC)
|
||||
return now > expires
|
||||
|
||||
@field_validator("code_challenge_method")
|
||||
def validate_challenge_method(cls, v):
|
||||
"""OAuth 2.1: Only S256 is allowed"""
|
||||
if v != "S256":
|
||||
raise ValueError("Only S256 code_challenge_method is supported (OAuth 2.1)")
|
||||
return v
|
||||
|
||||
class AccessToken(BaseModel):
|
||||
"""Access Token Model"""
|
||||
|
||||
token: str = Field(..., description="JWT access token")
|
||||
client_id: str
|
||||
scope: str
|
||||
expires_at: datetime
|
||||
user_id: str | None = None
|
||||
project_id: str = Field(default="*", description="Project ID for scoping")
|
||||
issued_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
def is_expired(self) -> bool:
|
||||
"""Check if token is expired"""
|
||||
now = datetime.now(UTC)
|
||||
expires = self.expires_at
|
||||
if expires.tzinfo is None:
|
||||
expires = expires.replace(tzinfo=UTC)
|
||||
return now > expires
|
||||
|
||||
class RefreshToken(BaseModel):
|
||||
"""Refresh Token Model"""
|
||||
|
||||
token: str = Field(..., description="Refresh token (secure random)")
|
||||
client_id: str
|
||||
access_token: str | None = None
|
||||
expires_at: datetime
|
||||
revoked: bool = Field(default=False)
|
||||
rotation_count: int = Field(default=0, description="Number of times rotated")
|
||||
issued_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
def is_expired(self) -> bool:
|
||||
"""Check if token is expired"""
|
||||
now = datetime.now(UTC)
|
||||
expires = self.expires_at
|
||||
if expires.tzinfo is None:
|
||||
expires = expires.replace(tzinfo=UTC)
|
||||
return now > expires
|
||||
|
||||
class TokenRequest(BaseModel):
|
||||
"""Token Request Model"""
|
||||
|
||||
grant_type: str = Field(
|
||||
..., description="authorization_code | refresh_token | client_credentials"
|
||||
)
|
||||
|
||||
# For authorization_code
|
||||
code: str | None = None
|
||||
redirect_uri: str | None = None
|
||||
code_verifier: str | None = None
|
||||
|
||||
# For refresh_token
|
||||
refresh_token: str | None = None
|
||||
|
||||
# Common
|
||||
client_id: str
|
||||
client_secret: str | None = None
|
||||
scope: str | None = None
|
||||
|
||||
class TokenResponse(BaseModel):
|
||||
"""Token Response Model"""
|
||||
|
||||
access_token: str
|
||||
token_type: str = "Bearer"
|
||||
expires_in: int
|
||||
refresh_token: str | None = None
|
||||
scope: str
|
||||
396
core/oauth/server.py
Normal file
396
core/oauth/server.py
Normal file
@@ -0,0 +1,396 @@
|
||||
"""
|
||||
OAuth 2.1 Authorization Server
|
||||
Handles OAuth flows: authorization_code, refresh_token, client_credentials
|
||||
"""
|
||||
|
||||
import logging
|
||||
import secrets
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
from .client_registry import get_client_registry
|
||||
from .pkce import validate_code_challenge
|
||||
from .schemas import AuthorizationCode, TokenResponse
|
||||
from .storage import get_storage
|
||||
from .token_manager import get_token_manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class OAuthError(Exception):
|
||||
"""OAuth error with error code and description"""
|
||||
|
||||
def __init__(self, error: str, error_description: str, status_code: int = 400):
|
||||
self.error = error
|
||||
self.error_description = error_description
|
||||
self.status_code = status_code
|
||||
super().__init__(error_description)
|
||||
|
||||
class OAuthServer:
|
||||
"""
|
||||
OAuth 2.1 Authorization Server
|
||||
|
||||
Implements:
|
||||
- Authorization Code Grant with PKCE (mandatory)
|
||||
- Refresh Token Grant with rotation
|
||||
- Client Credentials Grant (machine-to-machine)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.client_registry = get_client_registry()
|
||||
self.token_manager = get_token_manager()
|
||||
self.storage = get_storage()
|
||||
|
||||
# Authorization code TTL (5 minutes)
|
||||
self.auth_code_ttl = 300
|
||||
|
||||
def validate_authorization_request(
|
||||
self,
|
||||
client_id: str,
|
||||
redirect_uri: str,
|
||||
response_type: str,
|
||||
code_challenge: str,
|
||||
code_challenge_method: str,
|
||||
scope: str | None = None,
|
||||
state: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Validate OAuth authorization request (Step 1 of Authorization Code flow)
|
||||
|
||||
Args:
|
||||
client_id: OAuth client ID
|
||||
redirect_uri: Callback URI for authorization code
|
||||
response_type: Must be "code" (OAuth 2.1)
|
||||
code_challenge: PKCE code challenge
|
||||
code_challenge_method: Must be "S256" (OAuth 2.1)
|
||||
scope: Requested scopes (space-separated)
|
||||
state: Optional state parameter for CSRF protection
|
||||
|
||||
Returns:
|
||||
Dict with validated parameters
|
||||
|
||||
Raises:
|
||||
OAuthError: If validation fails
|
||||
"""
|
||||
# Validate client
|
||||
client = self.client_registry.get_client(client_id)
|
||||
if not client:
|
||||
raise OAuthError(
|
||||
error="invalid_client",
|
||||
error_description=f"Client {client_id} not found",
|
||||
status_code=401,
|
||||
)
|
||||
|
||||
# Validate response_type (OAuth 2.1: only "code" is allowed)
|
||||
if response_type != "code":
|
||||
raise OAuthError(
|
||||
error="unsupported_response_type",
|
||||
error_description="Only 'code' response_type is supported (OAuth 2.1)",
|
||||
)
|
||||
|
||||
if "authorization_code" not in client.grant_types:
|
||||
raise OAuthError(
|
||||
error="unauthorized_client",
|
||||
error_description="Client not authorized for authorization_code grant",
|
||||
)
|
||||
|
||||
# Validate redirect_uri (exact match)
|
||||
if redirect_uri not in client.redirect_uris:
|
||||
raise OAuthError(
|
||||
error="invalid_request", error_description=f"Invalid redirect_uri: {redirect_uri}"
|
||||
)
|
||||
|
||||
# Validate PKCE (mandatory in OAuth 2.1)
|
||||
if not code_challenge or not code_challenge_method:
|
||||
raise OAuthError(
|
||||
error="invalid_request",
|
||||
error_description="code_challenge and code_challenge_method are required (OAuth 2.1)",
|
||||
)
|
||||
|
||||
if code_challenge_method != "S256":
|
||||
raise OAuthError(
|
||||
error="invalid_request",
|
||||
error_description="Only S256 code_challenge_method is supported (OAuth 2.1)",
|
||||
)
|
||||
|
||||
# Validate scope
|
||||
requested_scopes = scope.split() if scope else ["read"]
|
||||
for s in requested_scopes:
|
||||
if s not in client.allowed_scopes:
|
||||
raise OAuthError(
|
||||
error="invalid_scope",
|
||||
error_description=f"Scope '{s}' not allowed for this client",
|
||||
)
|
||||
|
||||
return {
|
||||
"client_id": client_id,
|
||||
"redirect_uri": redirect_uri,
|
||||
"scope": " ".join(requested_scopes),
|
||||
"code_challenge": code_challenge,
|
||||
"code_challenge_method": code_challenge_method,
|
||||
"state": state,
|
||||
}
|
||||
|
||||
def create_authorization_code(
|
||||
self,
|
||||
client_id: str,
|
||||
redirect_uri: str,
|
||||
scope: str,
|
||||
code_challenge: str,
|
||||
code_challenge_method: str = "S256",
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
api_key_project_id: str | None = None,
|
||||
api_key_scope: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Create authorization code (Step 2 of Authorization Code flow)
|
||||
|
||||
Args:
|
||||
client_id: OAuth client ID
|
||||
redirect_uri: Redirect URI
|
||||
scope: Granted scopes
|
||||
code_challenge: PKCE code challenge
|
||||
code_challenge_method: PKCE method (S256)
|
||||
user_id: Optional user ID (for user-based auth)
|
||||
api_key_id: Optional API Key ID for scope/project inheritance
|
||||
api_key_project_id: Optional project ID from API Key
|
||||
api_key_scope: Optional scope from API Key
|
||||
|
||||
Returns:
|
||||
Authorization code (valid for 5 minutes)
|
||||
"""
|
||||
# Generate secure random code
|
||||
code = f"auth_{secrets.token_urlsafe(32)}"
|
||||
|
||||
# Create authorization code
|
||||
auth_code = AuthorizationCode(
|
||||
code=code,
|
||||
client_id=client_id,
|
||||
redirect_uri=redirect_uri,
|
||||
scope=scope,
|
||||
code_challenge=code_challenge,
|
||||
code_challenge_method=code_challenge_method,
|
||||
expires_at=datetime.now(UTC) + timedelta(seconds=self.auth_code_ttl),
|
||||
used=False,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
api_key_project_id=api_key_project_id,
|
||||
api_key_scope=api_key_scope,
|
||||
)
|
||||
|
||||
# Save to storage
|
||||
self.storage.save_authorization_code(auth_code)
|
||||
|
||||
logger.info(f"Created authorization code for client {client_id}")
|
||||
|
||||
return code
|
||||
|
||||
def exchange_code_for_tokens(
|
||||
self, client_id: str, client_secret: str, code: str, redirect_uri: str, code_verifier: str
|
||||
) -> TokenResponse:
|
||||
"""
|
||||
Exchange authorization code for tokens (Step 3 of Authorization Code flow)
|
||||
|
||||
Args:
|
||||
client_id: OAuth client ID
|
||||
client_secret: Client secret
|
||||
code: Authorization code from /authorize
|
||||
redirect_uri: Same redirect_uri used in /authorize
|
||||
code_verifier: PKCE code verifier
|
||||
|
||||
Returns:
|
||||
TokenResponse with access_token and refresh_token
|
||||
|
||||
Raises:
|
||||
OAuthError: If validation fails
|
||||
"""
|
||||
# Validate client credentials
|
||||
if not self.client_registry.validate_client_secret(client_id, client_secret):
|
||||
raise OAuthError(
|
||||
error="invalid_client",
|
||||
error_description="Invalid client credentials",
|
||||
status_code=401,
|
||||
)
|
||||
|
||||
# Get authorization code
|
||||
auth_code = self.storage.get_authorization_code(code)
|
||||
if not auth_code:
|
||||
raise OAuthError(
|
||||
error="invalid_grant", error_description="Invalid or expired authorization code"
|
||||
)
|
||||
|
||||
# Check if already used (prevents replay attacks)
|
||||
if auth_code.used:
|
||||
# Revoke all tokens for this client (security measure)
|
||||
logger.critical(
|
||||
f"Authorization code reuse detected for client {client_id}! "
|
||||
f"Code: {code[:20]}..."
|
||||
)
|
||||
raise OAuthError(
|
||||
error="invalid_grant", error_description="Authorization code already used"
|
||||
)
|
||||
|
||||
# Validate client_id match
|
||||
if auth_code.client_id != client_id:
|
||||
raise OAuthError(error="invalid_grant", error_description="Client ID mismatch")
|
||||
|
||||
# Validate redirect_uri match
|
||||
if auth_code.redirect_uri != redirect_uri:
|
||||
raise OAuthError(error="invalid_grant", error_description="Redirect URI mismatch")
|
||||
|
||||
# Validate PKCE code_verifier
|
||||
if not validate_code_challenge(
|
||||
code_verifier, auth_code.code_challenge, auth_code.code_challenge_method
|
||||
):
|
||||
raise OAuthError(
|
||||
error="invalid_grant",
|
||||
error_description="Invalid code_verifier (PKCE validation failed)",
|
||||
)
|
||||
|
||||
# Mark code as used
|
||||
auth_code.used = True
|
||||
self.storage.update_authorization_code(code, auth_code)
|
||||
|
||||
# Generate tokens with API Key's project and scope
|
||||
# If authorization code has API Key metadata, use it for scoping
|
||||
project_id = auth_code.api_key_project_id or "*"
|
||||
token_scope = auth_code.api_key_scope or auth_code.scope
|
||||
|
||||
access_token = self.token_manager.generate_access_token(
|
||||
client_id=client_id,
|
||||
scope=token_scope,
|
||||
user_id=auth_code.user_id or auth_code.api_key_id,
|
||||
project_id=project_id,
|
||||
)
|
||||
|
||||
refresh_token = self.token_manager.generate_refresh_token(
|
||||
client_id=client_id, access_token=access_token
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Exchanged authorization code for tokens: {client_id} "
|
||||
f"(project_id={project_id}, scope={token_scope})"
|
||||
)
|
||||
|
||||
return TokenResponse(
|
||||
access_token=access_token,
|
||||
token_type="Bearer",
|
||||
expires_in=self.token_manager.access_token_ttl,
|
||||
refresh_token=refresh_token,
|
||||
scope=auth_code.scope,
|
||||
)
|
||||
|
||||
def handle_refresh_token_grant(
|
||||
self, client_id: str, client_secret: str, refresh_token: str
|
||||
) -> TokenResponse:
|
||||
"""
|
||||
Handle refresh token grant (refresh access token)
|
||||
|
||||
Args:
|
||||
client_id: OAuth client ID
|
||||
client_secret: Client secret
|
||||
refresh_token: Current refresh token
|
||||
|
||||
Returns:
|
||||
TokenResponse with new access_token and refresh_token
|
||||
|
||||
Raises:
|
||||
OAuthError: If validation fails
|
||||
"""
|
||||
# Validate client credentials
|
||||
if not self.client_registry.validate_client_secret(client_id, client_secret):
|
||||
raise OAuthError(
|
||||
error="invalid_client",
|
||||
error_description="Invalid client credentials",
|
||||
status_code=401,
|
||||
)
|
||||
|
||||
# Check grant type is allowed
|
||||
client = self.client_registry.get_client(client_id)
|
||||
if "refresh_token" not in client.grant_types:
|
||||
raise OAuthError(
|
||||
error="unauthorized_client",
|
||||
error_description="Client not authorized for refresh_token grant",
|
||||
)
|
||||
|
||||
try:
|
||||
# Rotate refresh token
|
||||
new_tokens = self.token_manager.rotate_refresh_token(
|
||||
refresh_token=refresh_token, client_id=client_id
|
||||
)
|
||||
|
||||
return TokenResponse(**new_tokens)
|
||||
|
||||
except ValueError as e:
|
||||
raise OAuthError(error="invalid_grant", error_description=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error rotating refresh token: {e}")
|
||||
raise OAuthError(
|
||||
error="server_error", error_description="Internal server error", status_code=500
|
||||
)
|
||||
|
||||
def handle_client_credentials_grant(
|
||||
self, client_id: str, client_secret: str, scope: str | None = None
|
||||
) -> TokenResponse:
|
||||
"""
|
||||
Handle client credentials grant (machine-to-machine)
|
||||
|
||||
Args:
|
||||
client_id: OAuth client ID
|
||||
client_secret: Client secret
|
||||
scope: Requested scopes (space-separated)
|
||||
|
||||
Returns:
|
||||
TokenResponse with access_token (no refresh_token)
|
||||
|
||||
Raises:
|
||||
OAuthError: If validation fails
|
||||
"""
|
||||
# Validate client credentials
|
||||
if not self.client_registry.validate_client_secret(client_id, client_secret):
|
||||
raise OAuthError(
|
||||
error="invalid_client",
|
||||
error_description="Invalid client credentials",
|
||||
status_code=401,
|
||||
)
|
||||
|
||||
# Check grant type is allowed
|
||||
client = self.client_registry.get_client(client_id)
|
||||
if "client_credentials" not in client.grant_types:
|
||||
raise OAuthError(
|
||||
error="unauthorized_client",
|
||||
error_description="Client not authorized for client_credentials grant",
|
||||
)
|
||||
|
||||
# Validate scope
|
||||
requested_scopes = scope.split() if scope else [client.scope]
|
||||
for s in requested_scopes:
|
||||
if s not in client.allowed_scopes:
|
||||
raise OAuthError(
|
||||
error="invalid_scope",
|
||||
error_description=f"Scope '{s}' not allowed for this client",
|
||||
)
|
||||
|
||||
# Generate access token (no refresh token for client credentials)
|
||||
access_token = self.token_manager.generate_access_token(
|
||||
client_id=client_id, scope=" ".join(requested_scopes)
|
||||
)
|
||||
|
||||
logger.info(f"Generated client credentials token for {client_id}")
|
||||
|
||||
return TokenResponse(
|
||||
access_token=access_token,
|
||||
token_type="Bearer",
|
||||
expires_in=self.token_manager.access_token_ttl,
|
||||
scope=" ".join(requested_scopes),
|
||||
)
|
||||
|
||||
# Singleton
|
||||
_oauth_server: OAuthServer | None = None
|
||||
|
||||
def get_oauth_server() -> OAuthServer:
|
||||
"""Get singleton OAuthServer instance"""
|
||||
global _oauth_server
|
||||
if _oauth_server is None:
|
||||
_oauth_server = OAuthServer()
|
||||
return _oauth_server
|
||||
221
core/oauth/storage.py
Normal file
221
core/oauth/storage.py
Normal file
@@ -0,0 +1,221 @@
|
||||
"""
|
||||
OAuth 2.1 Token Storage
|
||||
Supports JSON files (default) and Redis (optional)
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .schemas import AccessToken, AuthorizationCode, RefreshToken
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class BaseStorage:
|
||||
"""Base storage interface"""
|
||||
|
||||
def save_authorization_code(self, code_data: AuthorizationCode) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_authorization_code(self, code: str) -> AuthorizationCode | None:
|
||||
raise NotImplementedError
|
||||
|
||||
def update_authorization_code(self, code: str, code_data: AuthorizationCode) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def delete_authorization_code(self, code: str) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def save_access_token(self, token_data: AccessToken) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_access_token(self, token: str) -> AccessToken | None:
|
||||
raise NotImplementedError
|
||||
|
||||
def save_refresh_token(self, token_data: RefreshToken) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_refresh_token(self, token: str) -> RefreshToken | None:
|
||||
raise NotImplementedError
|
||||
|
||||
def revoke_refresh_token(self, token: str) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
class JSONStorage(BaseStorage):
|
||||
"""
|
||||
JSON file storage for OAuth tokens.
|
||||
|
||||
Storage structure:
|
||||
data/oauth_codes.json - Authorization codes (short-lived)
|
||||
data/oauth_access_tokens.json - Access tokens
|
||||
data/oauth_refresh_tokens.json - Refresh tokens
|
||||
"""
|
||||
|
||||
def __init__(self, data_dir: str = "/app/data"):
|
||||
self.data_dir = Path(data_dir)
|
||||
self.data_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.codes_file = self.data_dir / "oauth_codes.json"
|
||||
self.access_tokens_file = self.data_dir / "oauth_access_tokens.json"
|
||||
self.refresh_tokens_file = self.data_dir / "oauth_refresh_tokens.json"
|
||||
|
||||
# Initialize files if not exist
|
||||
for file in [self.codes_file, self.access_tokens_file, self.refresh_tokens_file]:
|
||||
if not file.exists():
|
||||
self._write_json(file, {})
|
||||
|
||||
def _read_json(self, file_path: Path) -> dict[str, Any]:
|
||||
"""Read JSON file"""
|
||||
try:
|
||||
with open(file_path) as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
logger.error(f"Error reading {file_path}: {e}")
|
||||
return {}
|
||||
|
||||
def _write_json(self, file_path: Path, data: dict[str, Any]) -> bool:
|
||||
"""Write JSON file"""
|
||||
try:
|
||||
with open(file_path, "w") as f:
|
||||
json.dump(data, f, indent=2, default=str)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error writing {file_path}: {e}")
|
||||
return False
|
||||
|
||||
def save_authorization_code(self, code_data: AuthorizationCode) -> bool:
|
||||
"""Save authorization code"""
|
||||
codes = self._read_json(self.codes_file)
|
||||
codes[code_data.code] = code_data.model_dump(mode="json")
|
||||
return self._write_json(self.codes_file, codes)
|
||||
|
||||
def get_authorization_code(self, code: str) -> AuthorizationCode | None:
|
||||
"""Get authorization code"""
|
||||
codes = self._read_json(self.codes_file)
|
||||
code_data = codes.get(code)
|
||||
|
||||
if not code_data:
|
||||
return None
|
||||
|
||||
# Parse and check expiry
|
||||
auth_code = AuthorizationCode(**code_data)
|
||||
|
||||
if auth_code.is_expired():
|
||||
# Auto-cleanup expired codes
|
||||
self.delete_authorization_code(code)
|
||||
return None
|
||||
|
||||
return auth_code
|
||||
|
||||
def update_authorization_code(self, code: str, code_data: AuthorizationCode) -> bool:
|
||||
"""Update authorization code (e.g., mark as used)"""
|
||||
return self.save_authorization_code(code_data)
|
||||
|
||||
def delete_authorization_code(self, code: str) -> bool:
|
||||
"""Delete authorization code"""
|
||||
codes = self._read_json(self.codes_file)
|
||||
if code in codes:
|
||||
del codes[code]
|
||||
return self._write_json(self.codes_file, codes)
|
||||
return False
|
||||
|
||||
def save_access_token(self, token_data: AccessToken) -> bool:
|
||||
"""Save access token"""
|
||||
tokens = self._read_json(self.access_tokens_file)
|
||||
tokens[token_data.token] = token_data.model_dump(mode="json")
|
||||
return self._write_json(self.access_tokens_file, tokens)
|
||||
|
||||
def get_access_token(self, token: str) -> AccessToken | None:
|
||||
"""Get access token"""
|
||||
tokens = self._read_json(self.access_tokens_file)
|
||||
token_data = tokens.get(token)
|
||||
|
||||
if not token_data:
|
||||
return None
|
||||
|
||||
access_token = AccessToken(**token_data)
|
||||
|
||||
if access_token.is_expired():
|
||||
# Auto-cleanup
|
||||
tokens.pop(token, None)
|
||||
self._write_json(self.access_tokens_file, tokens)
|
||||
return None
|
||||
|
||||
return access_token
|
||||
|
||||
def save_refresh_token(self, token_data: RefreshToken) -> bool:
|
||||
"""Save refresh token"""
|
||||
tokens = self._read_json(self.refresh_tokens_file)
|
||||
tokens[token_data.token] = token_data.model_dump(mode="json")
|
||||
return self._write_json(self.refresh_tokens_file, tokens)
|
||||
|
||||
def get_refresh_token(self, token: str, include_revoked: bool = False) -> RefreshToken | None:
|
||||
"""Get refresh token.
|
||||
|
||||
Args:
|
||||
token: Refresh token string
|
||||
include_revoked: If True, return revoked tokens (for reuse detection)
|
||||
"""
|
||||
tokens = self._read_json(self.refresh_tokens_file)
|
||||
token_data = tokens.get(token)
|
||||
|
||||
if not token_data:
|
||||
return None
|
||||
|
||||
refresh_token = RefreshToken(**token_data)
|
||||
|
||||
if refresh_token.is_expired():
|
||||
return None
|
||||
|
||||
if refresh_token.revoked and not include_revoked:
|
||||
return None
|
||||
|
||||
return refresh_token
|
||||
|
||||
def revoke_refresh_token(self, token: str) -> bool:
|
||||
"""Revoke refresh token"""
|
||||
tokens = self._read_json(self.refresh_tokens_file)
|
||||
|
||||
if token in tokens:
|
||||
tokens[token]["revoked"] = True
|
||||
return self._write_json(self.refresh_tokens_file, tokens)
|
||||
|
||||
return False
|
||||
|
||||
def cleanup_expired(self):
|
||||
"""Cleanup expired tokens (run periodically)"""
|
||||
# Cleanup authorization codes
|
||||
codes = self._read_json(self.codes_file)
|
||||
cleaned_codes = {k: v for k, v in codes.items() if not AuthorizationCode(**v).is_expired()}
|
||||
self._write_json(self.codes_file, cleaned_codes)
|
||||
|
||||
# Cleanup access tokens
|
||||
tokens = self._read_json(self.access_tokens_file)
|
||||
cleaned_tokens = {k: v for k, v in tokens.items() if not AccessToken(**v).is_expired()}
|
||||
self._write_json(self.access_tokens_file, cleaned_tokens)
|
||||
|
||||
logger.info(f"Cleaned up {len(codes) - len(cleaned_codes)} expired authorization codes")
|
||||
logger.info(f"Cleaned up {len(tokens) - len(cleaned_tokens)} expired access tokens")
|
||||
|
||||
def get_storage() -> BaseStorage:
|
||||
"""
|
||||
Get storage instance based on environment.
|
||||
|
||||
Environment Variables:
|
||||
OAUTH_STORAGE_TYPE: "json" (default) | "redis"
|
||||
OAUTH_STORAGE_PATH: Path for JSON files (default: /app/data)
|
||||
"""
|
||||
storage_type = os.getenv("OAUTH_STORAGE_TYPE", "json")
|
||||
|
||||
if storage_type == "json":
|
||||
storage_path = os.getenv("OAUTH_STORAGE_PATH", "/app/data")
|
||||
return JSONStorage(storage_path)
|
||||
|
||||
# Future: Redis support
|
||||
# elif storage_type == "redis":
|
||||
# return RedisStorage()
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown storage type: {storage_type}")
|
||||
313
core/oauth/token_manager.py
Normal file
313
core/oauth/token_manager.py
Normal file
@@ -0,0 +1,313 @@
|
||||
"""
|
||||
OAuth 2.1 Token Manager
|
||||
JWT generation, validation, and refresh token rotation
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import secrets
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
import jwt
|
||||
|
||||
from .schemas import AccessToken, RefreshToken
|
||||
from .storage import get_storage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class SecurityError(Exception):
|
||||
"""Security-related error (e.g., token reuse)"""
|
||||
|
||||
pass
|
||||
|
||||
class TokenManager:
|
||||
"""
|
||||
OAuth 2.1 Token Manager.
|
||||
|
||||
Features:
|
||||
- JWT access tokens (1 hour default)
|
||||
- Refresh tokens with rotation (7 days default)
|
||||
- Refresh token reuse detection
|
||||
- Audit logging
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.storage = get_storage()
|
||||
|
||||
# JWT Configuration — require explicit secret or generate random fallback
|
||||
self.jwt_secret = os.getenv("OAUTH_JWT_SECRET_KEY")
|
||||
if not self.jwt_secret:
|
||||
self.jwt_secret = secrets.token_urlsafe(64)
|
||||
logger.warning(
|
||||
"OAUTH_JWT_SECRET_KEY not set. Generated random JWT secret. "
|
||||
"All tokens will be invalidated on restart. "
|
||||
"Set OAUTH_JWT_SECRET_KEY in your .env for persistent tokens."
|
||||
)
|
||||
self.jwt_algorithm = os.getenv("OAUTH_JWT_ALGORITHM", "HS256")
|
||||
|
||||
# Token TTLs (seconds)
|
||||
self.access_token_ttl = int(os.getenv("OAUTH_ACCESS_TOKEN_TTL", "3600")) # 1 hour
|
||||
self.refresh_token_ttl = int(os.getenv("OAUTH_REFRESH_TOKEN_TTL", "604800")) # 7 days
|
||||
|
||||
def generate_access_token(
|
||||
self, client_id: str, scope: str, user_id: str | None = None, project_id: str = "*"
|
||||
) -> str:
|
||||
"""
|
||||
Generate JWT access token.
|
||||
|
||||
Args:
|
||||
client_id: OAuth client ID
|
||||
scope: Granted scopes (space-separated)
|
||||
user_id: User ID (optional, for user-based auth)
|
||||
project_id: Project ID for scoping (default: "*" for global)
|
||||
|
||||
Returns:
|
||||
JWT access token
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
now_ts = int(time.time())
|
||||
exp_ts = now_ts + self.access_token_ttl
|
||||
expires_at = now + timedelta(seconds=self.access_token_ttl)
|
||||
|
||||
# JWT payload (use time.time() for correct UTC timestamps)
|
||||
payload = {
|
||||
"client_id": client_id,
|
||||
"scope": scope,
|
||||
"project_id": project_id,
|
||||
"iat": now_ts,
|
||||
"exp": exp_ts,
|
||||
"nbf": now_ts, # Not before
|
||||
"jti": secrets.token_urlsafe(16), # JWT ID (unique)
|
||||
}
|
||||
|
||||
if user_id:
|
||||
payload["sub"] = user_id # Subject (user ID)
|
||||
|
||||
# Encode JWT
|
||||
token = jwt.encode(payload, self.jwt_secret, algorithm=self.jwt_algorithm)
|
||||
|
||||
# Store token metadata
|
||||
token_data = AccessToken(
|
||||
token=token,
|
||||
client_id=client_id,
|
||||
scope=scope,
|
||||
expires_at=expires_at,
|
||||
user_id=user_id,
|
||||
project_id=project_id,
|
||||
issued_at=now,
|
||||
)
|
||||
self.storage.save_access_token(token_data)
|
||||
|
||||
logger.info(f"Generated access token for client {client_id} (scope: {scope})")
|
||||
|
||||
return token
|
||||
|
||||
def validate_access_token(self, token: str) -> dict[str, Any]:
|
||||
"""
|
||||
Validate JWT access token.
|
||||
|
||||
Args:
|
||||
token: JWT access token
|
||||
|
||||
Returns:
|
||||
Decoded payload
|
||||
|
||||
Raises:
|
||||
jwt.ExpiredSignatureError: Token expired
|
||||
jwt.InvalidTokenError: Invalid token
|
||||
"""
|
||||
try:
|
||||
# Decode and verify JWT
|
||||
payload = jwt.decode(
|
||||
token,
|
||||
self.jwt_secret,
|
||||
algorithms=[self.jwt_algorithm],
|
||||
options={
|
||||
"verify_signature": True,
|
||||
"verify_exp": True,
|
||||
"verify_nbf": True,
|
||||
},
|
||||
)
|
||||
|
||||
return payload
|
||||
|
||||
except jwt.ExpiredSignatureError:
|
||||
logger.warning("Expired access token")
|
||||
raise
|
||||
except jwt.InvalidTokenError as e:
|
||||
logger.warning(f"Invalid access token: {e}")
|
||||
raise
|
||||
|
||||
def generate_refresh_token(self, client_id: str, access_token: str) -> str:
|
||||
"""
|
||||
Generate refresh token.
|
||||
|
||||
Args:
|
||||
client_id: OAuth client ID
|
||||
access_token: Associated access token
|
||||
|
||||
Returns:
|
||||
Refresh token (secure random string)
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
expires_at = now + timedelta(seconds=self.refresh_token_ttl)
|
||||
|
||||
# Generate secure random token
|
||||
token = f"rt_{secrets.token_urlsafe(32)}"
|
||||
|
||||
# Store refresh token
|
||||
token_data = RefreshToken(
|
||||
token=token,
|
||||
client_id=client_id,
|
||||
access_token=access_token,
|
||||
expires_at=expires_at,
|
||||
revoked=False,
|
||||
rotation_count=0,
|
||||
issued_at=now,
|
||||
)
|
||||
self.storage.save_refresh_token(token_data)
|
||||
|
||||
logger.info(f"Generated refresh token for client {client_id}")
|
||||
|
||||
return token
|
||||
|
||||
def rotate_refresh_token(self, refresh_token: str, client_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
Rotate refresh token (OAuth 2.1 best practice).
|
||||
|
||||
Security:
|
||||
- Old refresh token is immediately revoked
|
||||
- Reuse detection: if old token is used again, revoke all tokens
|
||||
|
||||
Args:
|
||||
refresh_token: Current refresh token
|
||||
client_id: OAuth client ID
|
||||
|
||||
Returns:
|
||||
{
|
||||
"access_token": str,
|
||||
"refresh_token": str,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": int,
|
||||
"scope": str
|
||||
}
|
||||
|
||||
Raises:
|
||||
ValueError: Invalid/expired refresh token
|
||||
SecurityError: Refresh token reuse detected
|
||||
"""
|
||||
# Get refresh token (include revoked for reuse detection)
|
||||
token_data = self.storage.get_refresh_token(refresh_token, include_revoked=True)
|
||||
|
||||
if not token_data:
|
||||
raise ValueError("Invalid or expired refresh token")
|
||||
|
||||
# Check if revoked (reuse detection!)
|
||||
if token_data.revoked:
|
||||
logger.critical(
|
||||
f"Refresh token reuse detected for client {client_id}! "
|
||||
f"Revoking all tokens for this client."
|
||||
)
|
||||
|
||||
# Log security event
|
||||
try:
|
||||
from core.audit_log import LogLevel, get_audit_logger
|
||||
|
||||
audit_logger = get_audit_logger()
|
||||
audit_logger.log_system_event(
|
||||
event=f"SECURITY: Refresh token reuse detected: {client_id}",
|
||||
details={"client_id": client_id, "token": refresh_token[:20] + "..."},
|
||||
level=LogLevel.CRITICAL,
|
||||
)
|
||||
except (ImportError, Exception):
|
||||
# Audit logging not available or failed
|
||||
pass
|
||||
|
||||
# Revoke all tokens for this client (TODO: implement)
|
||||
# For now, just raise error
|
||||
raise SecurityError("Refresh token reuse detected - all tokens revoked")
|
||||
|
||||
# Validate client_id match
|
||||
if token_data.client_id != client_id:
|
||||
raise ValueError("Client ID mismatch")
|
||||
|
||||
# Get scope from old access token (if available)
|
||||
scope = "read" # Default
|
||||
if token_data.access_token:
|
||||
try:
|
||||
old_payload = self.validate_access_token(token_data.access_token)
|
||||
scope = old_payload.get("scope", "read")
|
||||
except:
|
||||
pass # Old token might be expired, use default
|
||||
|
||||
# Generate new tokens
|
||||
new_access_token = self.generate_access_token(
|
||||
client_id=client_id,
|
||||
scope=scope,
|
||||
user_id=None, # TODO: get from old token
|
||||
project_id="*",
|
||||
)
|
||||
|
||||
new_refresh_token = self.generate_refresh_token(
|
||||
client_id=client_id, access_token=new_access_token
|
||||
)
|
||||
|
||||
# Revoke old refresh token
|
||||
self.storage.revoke_refresh_token(refresh_token)
|
||||
|
||||
# Increment rotation count
|
||||
token_data.rotation_count += 1
|
||||
|
||||
logger.info(
|
||||
f"Rotated refresh token for client {client_id} "
|
||||
f"(rotation #{token_data.rotation_count})"
|
||||
)
|
||||
|
||||
# Log audit event
|
||||
try:
|
||||
from core.audit_log import LogLevel, get_audit_logger
|
||||
|
||||
audit_logger = get_audit_logger()
|
||||
audit_logger.log_system_event(
|
||||
event=f"Refresh token rotated for {client_id}",
|
||||
details={"client_id": client_id, "rotation_count": token_data.rotation_count},
|
||||
level=LogLevel.INFO,
|
||||
)
|
||||
except (ImportError, Exception):
|
||||
# Audit logging not available or failed
|
||||
pass
|
||||
|
||||
return {
|
||||
"access_token": new_access_token,
|
||||
"refresh_token": new_refresh_token,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": self.access_token_ttl,
|
||||
"scope": scope,
|
||||
}
|
||||
|
||||
def revoke_token(self, token: str, token_type: str = "refresh"):
|
||||
"""
|
||||
Revoke token.
|
||||
|
||||
Args:
|
||||
token: Token to revoke
|
||||
token_type: "refresh" or "access"
|
||||
"""
|
||||
if token_type == "refresh":
|
||||
self.storage.revoke_refresh_token(token)
|
||||
logger.info("Revoked refresh token")
|
||||
|
||||
# Access tokens cannot be revoked (JWT stateless)
|
||||
# They expire naturally after TTL
|
||||
|
||||
# Singleton
|
||||
_token_manager: TokenManager | None = None
|
||||
|
||||
def get_token_manager() -> TokenManager:
|
||||
"""Get singleton TokenManager instance"""
|
||||
global _token_manager
|
||||
if _token_manager is None:
|
||||
_token_manager = TokenManager()
|
||||
return _token_manager
|
||||
245
core/project_manager.py
Normal file
245
core/project_manager.py
Normal file
@@ -0,0 +1,245 @@
|
||||
"""
|
||||
Project Manager
|
||||
|
||||
Discovers and manages project instances from environment variables.
|
||||
Handles plugin lifecycle and tool registration.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from plugins import BasePlugin, registry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ProjectManager:
|
||||
"""
|
||||
Manage multiple project instances.
|
||||
|
||||
Projects are discovered from environment variables:
|
||||
- {PLUGIN_TYPE}_{PROJECT_ID}_{CONFIG_KEY}
|
||||
|
||||
Example:
|
||||
WORDPRESS_SITE1_URL=https://example.com
|
||||
WORDPRESS_SITE1_USERNAME=admin
|
||||
WORDPRESS_SITE1_APP_PASSWORD=xxxx
|
||||
WORDPRESS_SITE2_URL=https://other.com
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize project manager."""
|
||||
self.projects: dict[str, BasePlugin] = {}
|
||||
self.logger = logging.getLogger("ProjectManager")
|
||||
|
||||
def discover_projects(self) -> None:
|
||||
"""
|
||||
Discover projects from environment variables.
|
||||
|
||||
Scans environment for project configurations and creates
|
||||
plugin instances.
|
||||
"""
|
||||
self.logger.info("Starting project discovery...")
|
||||
|
||||
# Get all registered plugin types
|
||||
plugin_types = registry.get_registered_types()
|
||||
|
||||
for plugin_type in plugin_types:
|
||||
self._discover_plugin_type(plugin_type)
|
||||
|
||||
self.logger.info(f"Discovery complete. Found {len(self.projects)} projects.")
|
||||
|
||||
def _discover_plugin_type(self, plugin_type: str) -> None:
|
||||
"""
|
||||
Discover all projects of a specific plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Type of plugin (e.g., 'wordpress')
|
||||
"""
|
||||
prefix = plugin_type.upper() + "_"
|
||||
|
||||
# Find all project IDs for this plugin type
|
||||
project_ids = set()
|
||||
env_pattern = re.compile(f"^{prefix}([A-Z0-9_]+?)_(.+)$")
|
||||
|
||||
for env_key in os.environ.keys():
|
||||
match = env_pattern.match(env_key)
|
||||
if match:
|
||||
project_id = match.group(1).lower()
|
||||
project_ids.add(project_id)
|
||||
|
||||
# Create plugin instance for each project
|
||||
for project_id in project_ids:
|
||||
try:
|
||||
config = self._load_project_config(plugin_type, project_id)
|
||||
if config:
|
||||
self._create_project_instance(plugin_type, project_id, config)
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to create {plugin_type} project '{project_id}': {e}", exc_info=True
|
||||
)
|
||||
|
||||
def _load_project_config(self, plugin_type: str, project_id: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Load configuration for a project from environment.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type
|
||||
project_id: Project ID
|
||||
|
||||
Returns:
|
||||
Dict with configuration or None if incomplete
|
||||
"""
|
||||
prefix = f"{plugin_type.upper()}_{project_id.upper()}_"
|
||||
config = {}
|
||||
|
||||
# Collect all config keys for this project
|
||||
for env_key, env_value in os.environ.items():
|
||||
if env_key.startswith(prefix):
|
||||
# Extract config key (everything after prefix)
|
||||
config_key = env_key[len(prefix) :].lower()
|
||||
config[config_key] = env_value
|
||||
|
||||
if not config:
|
||||
return None
|
||||
|
||||
self.logger.debug(f"Loaded config for {plugin_type}/{project_id}: {list(config.keys())}")
|
||||
return config
|
||||
|
||||
def _create_project_instance(
|
||||
self, plugin_type: str, project_id: str, config: dict[str, Any]
|
||||
) -> None:
|
||||
"""
|
||||
Create a plugin instance for a project.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type
|
||||
project_id: Project ID
|
||||
config: Project configuration
|
||||
"""
|
||||
try:
|
||||
# Create plugin instance
|
||||
plugin = registry.create_instance(plugin_type, project_id, config)
|
||||
|
||||
# Store with full identifier
|
||||
full_id = f"{plugin_type}_{project_id}"
|
||||
self.projects[full_id] = plugin
|
||||
|
||||
self.logger.info(f"Created project: {full_id}")
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to instantiate {plugin_type}/{project_id}: {e}")
|
||||
|
||||
def get_project(self, full_id: str) -> BasePlugin | None:
|
||||
"""
|
||||
Get a project plugin instance.
|
||||
|
||||
Args:
|
||||
full_id: Full project identifier (plugin_type_project_id)
|
||||
|
||||
Returns:
|
||||
Plugin instance or None
|
||||
"""
|
||||
return self.projects.get(full_id)
|
||||
|
||||
def get_all_projects(self) -> dict[str, BasePlugin]:
|
||||
"""Get all project instances."""
|
||||
return self.projects.copy()
|
||||
|
||||
def get_projects_by_type(self, plugin_type: str) -> dict[str, BasePlugin]:
|
||||
"""
|
||||
Get all projects of a specific type.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type to filter by
|
||||
|
||||
Returns:
|
||||
Dict of project_id -> plugin
|
||||
"""
|
||||
prefix = plugin_type + "_"
|
||||
return {
|
||||
full_id: plugin
|
||||
for full_id, plugin in self.projects.items()
|
||||
if full_id.startswith(prefix)
|
||||
}
|
||||
|
||||
def get_all_tools(self) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Get all MCP tools from all projects.
|
||||
|
||||
Returns:
|
||||
List of tool definitions
|
||||
"""
|
||||
all_tools = []
|
||||
|
||||
for full_id, plugin in self.projects.items():
|
||||
try:
|
||||
tools = plugin.get_tools()
|
||||
all_tools.extend(tools)
|
||||
self.logger.debug(f"Loaded {len(tools)} tools from {full_id}")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error loading tools from {full_id}: {e}", exc_info=True)
|
||||
|
||||
self.logger.debug(f"Total tools loaded: {len(all_tools)}")
|
||||
return all_tools
|
||||
|
||||
async def check_all_health(self) -> dict[str, dict[str, Any]]:
|
||||
"""
|
||||
Check health of all projects.
|
||||
|
||||
Returns:
|
||||
Dict mapping project ID to health status
|
||||
"""
|
||||
health_results = {}
|
||||
|
||||
for full_id, plugin in self.projects.items():
|
||||
try:
|
||||
health = await plugin.health_check()
|
||||
health_results[full_id] = health
|
||||
except Exception as e:
|
||||
health_results[full_id] = {
|
||||
"healthy": False,
|
||||
"message": f"Health check failed: {str(e)}",
|
||||
}
|
||||
|
||||
return health_results
|
||||
|
||||
def get_project_info(self, full_id: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get information about a specific project.
|
||||
|
||||
Args:
|
||||
full_id: Full project identifier
|
||||
|
||||
Returns:
|
||||
Project info dict or None
|
||||
"""
|
||||
plugin = self.get_project(full_id)
|
||||
if plugin:
|
||||
return plugin.get_project_info()
|
||||
return None
|
||||
|
||||
def list_projects(self) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List all projects with basic information.
|
||||
|
||||
Returns:
|
||||
List of project info dicts
|
||||
"""
|
||||
return [
|
||||
{"id": full_id, "type": plugin.get_plugin_name(), "project_id": plugin.project_id}
|
||||
for full_id, plugin in self.projects.items()
|
||||
]
|
||||
|
||||
# Global project manager instance
|
||||
_project_manager: ProjectManager | None = None
|
||||
|
||||
def get_project_manager() -> ProjectManager:
|
||||
"""Get the global project manager instance."""
|
||||
global _project_manager
|
||||
if _project_manager is None:
|
||||
_project_manager = ProjectManager()
|
||||
_project_manager.discover_projects()
|
||||
return _project_manager
|
||||
414
core/rate_limiter.py
Normal file
414
core/rate_limiter.py
Normal file
@@ -0,0 +1,414 @@
|
||||
"""
|
||||
Rate Limiting & Throttling for MCP Server (Phase 7.3)
|
||||
|
||||
This module implements Token Bucket-based rate limiting to prevent API abuse
|
||||
and ensure fair resource usage across all MCP clients.
|
||||
|
||||
Features:
|
||||
- Multi-level rate limits (per minute, hour, day)
|
||||
- Per-client tracking with token bucket algorithm
|
||||
- Configurable limits per plugin type
|
||||
- Statistics and monitoring capabilities
|
||||
- Integration with audit logging
|
||||
|
||||
Author: Phase 7.3 Implementation
|
||||
Date: 2025-01-11
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class RateLimitConfig:
|
||||
"""Configuration for rate limits at different time intervals."""
|
||||
|
||||
per_minute: int = 60
|
||||
per_hour: int = 1000
|
||||
per_day: int = 10000
|
||||
|
||||
@classmethod
|
||||
def from_env(cls, prefix: str = "") -> "RateLimitConfig":
|
||||
"""Create config from environment variables."""
|
||||
env_prefix = f"{prefix}_" if prefix else ""
|
||||
return cls(
|
||||
per_minute=int(os.getenv(f"{env_prefix}RATE_LIMIT_PER_MINUTE", "60")),
|
||||
per_hour=int(os.getenv(f"{env_prefix}RATE_LIMIT_PER_HOUR", "1000")),
|
||||
per_day=int(os.getenv(f"{env_prefix}RATE_LIMIT_PER_DAY", "10000")),
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class TokenBucket:
|
||||
"""
|
||||
Token Bucket implementation for rate limiting.
|
||||
|
||||
The token bucket algorithm allows for burst traffic while maintaining
|
||||
an average rate limit over time.
|
||||
"""
|
||||
|
||||
capacity: int
|
||||
refill_rate: float # tokens per second
|
||||
tokens: float = field(default=0.0)
|
||||
last_refill: float = field(default_factory=time.time)
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize bucket with full capacity."""
|
||||
if self.tokens == 0.0:
|
||||
self.tokens = float(self.capacity)
|
||||
|
||||
def refill(self) -> None:
|
||||
"""Refill tokens based on elapsed time."""
|
||||
now = time.time()
|
||||
elapsed = now - self.last_refill
|
||||
|
||||
# Add tokens based on elapsed time
|
||||
self.tokens = min(self.capacity, self.tokens + (elapsed * self.refill_rate))
|
||||
self.last_refill = now
|
||||
|
||||
def consume(self, tokens: int = 1) -> bool:
|
||||
"""
|
||||
Attempt to consume tokens from the bucket.
|
||||
|
||||
Args:
|
||||
tokens: Number of tokens to consume
|
||||
|
||||
Returns:
|
||||
True if tokens were available and consumed, False otherwise
|
||||
"""
|
||||
self.refill()
|
||||
|
||||
if self.tokens >= tokens:
|
||||
self.tokens -= tokens
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_available_tokens(self) -> int:
|
||||
"""Get current number of available tokens."""
|
||||
self.refill()
|
||||
return int(self.tokens)
|
||||
|
||||
def get_wait_time(self, tokens: int = 1) -> float:
|
||||
"""
|
||||
Calculate wait time in seconds until enough tokens are available.
|
||||
|
||||
Args:
|
||||
tokens: Number of tokens needed
|
||||
|
||||
Returns:
|
||||
Wait time in seconds (0 if tokens already available)
|
||||
"""
|
||||
self.refill()
|
||||
|
||||
if self.tokens >= tokens:
|
||||
return 0.0
|
||||
|
||||
tokens_needed = tokens - self.tokens
|
||||
return tokens_needed / self.refill_rate
|
||||
|
||||
@dataclass
|
||||
class ClientRateLimitState:
|
||||
"""Track rate limit state for a single client."""
|
||||
|
||||
client_id: str
|
||||
minute_bucket: TokenBucket
|
||||
hour_bucket: TokenBucket
|
||||
day_bucket: TokenBucket
|
||||
total_requests: int = 0
|
||||
rejected_requests: int = 0
|
||||
last_request_time: float = field(default_factory=time.time)
|
||||
first_request_time: float = field(default_factory=time.time)
|
||||
|
||||
def check_and_consume(self) -> tuple[bool, str, float]:
|
||||
"""
|
||||
Check if request is allowed and consume tokens.
|
||||
|
||||
Returns:
|
||||
Tuple of (allowed, reason, retry_after_seconds)
|
||||
"""
|
||||
# Check each time window (most restrictive first)
|
||||
if not self.minute_bucket.consume():
|
||||
wait_time = self.minute_bucket.get_wait_time()
|
||||
self.rejected_requests += 1
|
||||
return False, "Rate limit exceeded: too many requests per minute", wait_time
|
||||
|
||||
if not self.hour_bucket.consume():
|
||||
wait_time = self.hour_bucket.get_wait_time()
|
||||
# Refund the minute token since we're rejecting
|
||||
self.minute_bucket.tokens = min(
|
||||
self.minute_bucket.capacity, self.minute_bucket.tokens + 1
|
||||
)
|
||||
self.rejected_requests += 1
|
||||
return False, "Rate limit exceeded: too many requests per hour", wait_time
|
||||
|
||||
if not self.day_bucket.consume():
|
||||
wait_time = self.day_bucket.get_wait_time()
|
||||
# Refund tokens since we're rejecting
|
||||
self.minute_bucket.tokens = min(
|
||||
self.minute_bucket.capacity, self.minute_bucket.tokens + 1
|
||||
)
|
||||
self.hour_bucket.tokens = min(self.hour_bucket.capacity, self.hour_bucket.tokens + 1)
|
||||
self.rejected_requests += 1
|
||||
return False, "Rate limit exceeded: daily limit reached", wait_time
|
||||
|
||||
# All checks passed
|
||||
self.total_requests += 1
|
||||
self.last_request_time = time.time()
|
||||
return True, "", 0.0
|
||||
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
"""Get statistics for this client."""
|
||||
now = time.time()
|
||||
uptime = now - self.first_request_time
|
||||
|
||||
return {
|
||||
"client_id": self.client_id,
|
||||
"total_requests": self.total_requests,
|
||||
"rejected_requests": self.rejected_requests,
|
||||
"success_rate": (
|
||||
(self.total_requests - self.rejected_requests) / self.total_requests
|
||||
if self.total_requests > 0
|
||||
else 1.0
|
||||
),
|
||||
"available_tokens": {
|
||||
"per_minute": self.minute_bucket.get_available_tokens(),
|
||||
"per_hour": self.hour_bucket.get_available_tokens(),
|
||||
"per_day": self.day_bucket.get_available_tokens(),
|
||||
},
|
||||
"limits": {
|
||||
"per_minute": self.minute_bucket.capacity,
|
||||
"per_hour": self.hour_bucket.capacity,
|
||||
"per_day": self.day_bucket.capacity,
|
||||
},
|
||||
"last_request": datetime.fromtimestamp(self.last_request_time, tz=UTC).isoformat(),
|
||||
"uptime_seconds": uptime,
|
||||
}
|
||||
|
||||
class RateLimiter:
|
||||
"""
|
||||
Rate limiter using Token Bucket algorithm.
|
||||
|
||||
Provides multi-level rate limiting (per minute, hour, day) with
|
||||
per-client tracking and configurable limits.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize rate limiter with default configuration."""
|
||||
self.clients: dict[str, ClientRateLimitState] = {}
|
||||
self.global_stats = {"total_requests": 0, "total_rejected": 0, "start_time": time.time()}
|
||||
|
||||
# Load default configuration from environment
|
||||
self.default_config = RateLimitConfig.from_env()
|
||||
|
||||
# Plugin-specific configurations
|
||||
self.plugin_configs: dict[str, RateLimitConfig] = {
|
||||
"wordpress": RateLimitConfig.from_env("WORDPRESS"),
|
||||
"woocommerce": RateLimitConfig.from_env("WOOCOMMERCE"),
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Rate limiter initialized with default limits: "
|
||||
f"{self.default_config.per_minute}/min, "
|
||||
f"{self.default_config.per_hour}/hour, "
|
||||
f"{self.default_config.per_day}/day"
|
||||
)
|
||||
|
||||
def _get_or_create_client_state(
|
||||
self, client_id: str, plugin_type: str | None = None
|
||||
) -> ClientRateLimitState:
|
||||
"""Get or create rate limit state for a client."""
|
||||
if client_id not in self.clients:
|
||||
# Determine which config to use
|
||||
config = self.plugin_configs.get(plugin_type, self.default_config)
|
||||
|
||||
# Create token buckets for each time window
|
||||
minute_bucket = TokenBucket(
|
||||
capacity=config.per_minute,
|
||||
refill_rate=config.per_minute / 60.0, # tokens per second
|
||||
)
|
||||
hour_bucket = TokenBucket(
|
||||
capacity=config.per_hour, refill_rate=config.per_hour / 3600.0
|
||||
)
|
||||
day_bucket = TokenBucket(capacity=config.per_day, refill_rate=config.per_day / 86400.0)
|
||||
|
||||
self.clients[client_id] = ClientRateLimitState(
|
||||
client_id=client_id,
|
||||
minute_bucket=minute_bucket,
|
||||
hour_bucket=hour_bucket,
|
||||
day_bucket=day_bucket,
|
||||
)
|
||||
|
||||
logger.debug(f"Created rate limit state for client: {client_id}")
|
||||
|
||||
return self.clients[client_id]
|
||||
|
||||
def check_rate_limit(
|
||||
self, client_id: str, tool_name: str | None = None, plugin_type: str | None = None
|
||||
) -> tuple[bool, str, float]:
|
||||
"""
|
||||
Check if request should be allowed based on rate limits.
|
||||
|
||||
Args:
|
||||
client_id: Identifier for the client (e.g., auth token hash)
|
||||
tool_name: Name of the tool being called (for logging)
|
||||
plugin_type: Type of plugin (wordpress, woocommerce, etc.)
|
||||
|
||||
Returns:
|
||||
Tuple of (allowed, message, retry_after_seconds)
|
||||
"""
|
||||
# Get or create client state
|
||||
client_state = self._get_or_create_client_state(client_id, plugin_type)
|
||||
|
||||
# Update global stats
|
||||
self.global_stats["total_requests"] += 1
|
||||
|
||||
# Check and consume tokens
|
||||
allowed, message, retry_after = client_state.check_and_consume()
|
||||
|
||||
if not allowed:
|
||||
# Track rejection
|
||||
client_state.rejected_requests += 1
|
||||
self.global_stats["total_rejected"] += 1
|
||||
|
||||
logger.warning(
|
||||
f"Rate limit exceeded for client {client_id[:8]}... "
|
||||
f"(tool: {tool_name}, reason: {message}, "
|
||||
f"retry_after: {retry_after:.1f}s)"
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
f"Rate limit check passed for client {client_id[:8]}... " f"(tool: {tool_name})"
|
||||
)
|
||||
|
||||
return allowed, message, retry_after
|
||||
|
||||
def get_client_stats(self, client_id: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get statistics for a specific client.
|
||||
|
||||
Args:
|
||||
client_id: Client identifier
|
||||
|
||||
Returns:
|
||||
Client statistics or None if client not found
|
||||
"""
|
||||
if client_id not in self.clients:
|
||||
return None
|
||||
|
||||
return self.clients[client_id].get_stats()
|
||||
|
||||
def get_all_stats(self) -> dict[str, Any]:
|
||||
"""Get global rate limiter statistics."""
|
||||
now = time.time()
|
||||
uptime = now - self.global_stats["start_time"]
|
||||
|
||||
# Calculate per-client stats
|
||||
client_stats = []
|
||||
for _client_id, client_state in self.clients.items():
|
||||
client_stats.append(client_state.get_stats())
|
||||
|
||||
return {
|
||||
"global": {
|
||||
"total_requests": self.global_stats["total_requests"],
|
||||
"total_rejected": self.global_stats["total_rejected"],
|
||||
"rejection_rate": (
|
||||
self.global_stats["total_rejected"] / self.global_stats["total_requests"]
|
||||
if self.global_stats["total_requests"] > 0
|
||||
else 0.0
|
||||
),
|
||||
"active_clients": len(self.clients),
|
||||
"uptime_seconds": uptime,
|
||||
"start_time": datetime.fromtimestamp(
|
||||
self.global_stats["start_time"], tz=UTC
|
||||
).isoformat(),
|
||||
},
|
||||
"default_limits": {
|
||||
"per_minute": self.default_config.per_minute,
|
||||
"per_hour": self.default_config.per_hour,
|
||||
"per_day": self.default_config.per_day,
|
||||
},
|
||||
"plugin_limits": {
|
||||
plugin: {
|
||||
"per_minute": config.per_minute,
|
||||
"per_hour": config.per_hour,
|
||||
"per_day": config.per_day,
|
||||
}
|
||||
for plugin, config in self.plugin_configs.items()
|
||||
},
|
||||
"clients": client_stats,
|
||||
}
|
||||
|
||||
def reset_client(self, client_id: str) -> bool:
|
||||
"""
|
||||
Reset rate limit state for a specific client.
|
||||
|
||||
Args:
|
||||
client_id: Client identifier
|
||||
|
||||
Returns:
|
||||
True if client was reset, False if client not found
|
||||
"""
|
||||
if client_id in self.clients:
|
||||
del self.clients[client_id]
|
||||
logger.info(f"Reset rate limit state for client: {client_id}")
|
||||
return True
|
||||
return False
|
||||
|
||||
def reset_all(self) -> int:
|
||||
"""
|
||||
Reset all client rate limit states.
|
||||
|
||||
Returns:
|
||||
Number of clients reset
|
||||
"""
|
||||
count = len(self.clients)
|
||||
self.clients.clear()
|
||||
self.global_stats = {"total_requests": 0, "total_rejected": 0, "start_time": time.time()}
|
||||
logger.info(f"Reset rate limit state for {count} clients")
|
||||
return count
|
||||
|
||||
def configure_limits(
|
||||
self,
|
||||
plugin_type: str,
|
||||
per_minute: int | None = None,
|
||||
per_hour: int | None = None,
|
||||
per_day: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Configure rate limits for a specific plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type identifier
|
||||
per_minute: Requests per minute limit
|
||||
per_hour: Requests per hour limit
|
||||
per_day: Requests per day limit
|
||||
"""
|
||||
if plugin_type not in self.plugin_configs:
|
||||
self.plugin_configs[plugin_type] = RateLimitConfig()
|
||||
|
||||
config = self.plugin_configs[plugin_type]
|
||||
if per_minute is not None:
|
||||
config.per_minute = per_minute
|
||||
if per_hour is not None:
|
||||
config.per_hour = per_hour
|
||||
if per_day is not None:
|
||||
config.per_day = per_day
|
||||
|
||||
logger.info(
|
||||
f"Updated rate limits for {plugin_type}: "
|
||||
f"{config.per_minute}/min, {config.per_hour}/hour, {config.per_day}/day"
|
||||
)
|
||||
|
||||
# Singleton instance
|
||||
_rate_limiter: RateLimiter | None = None
|
||||
|
||||
def get_rate_limiter() -> RateLimiter:
|
||||
"""Get or create the global rate limiter instance."""
|
||||
global _rate_limiter
|
||||
if _rate_limiter is None:
|
||||
_rate_limiter = RateLimiter()
|
||||
return _rate_limiter
|
||||
539
core/site_manager.py
Normal file
539
core/site_manager.py
Normal file
@@ -0,0 +1,539 @@
|
||||
"""
|
||||
Site Manager - Type-safe site configuration management
|
||||
|
||||
Manages site configurations with Pydantic validation.
|
||||
Part of Option B clean architecture refactoring.
|
||||
|
||||
Discovers sites from environment variables:
|
||||
- {PLUGIN_TYPE}_{SITE_ID}_{CONFIG_KEY}
|
||||
- {PLUGIN_TYPE}_{SITE_ID}_ALIAS (optional)
|
||||
|
||||
Example:
|
||||
WORDPRESS_SITE1_URL=https://example.com
|
||||
WORDPRESS_SITE1_USERNAME=admin
|
||||
WORDPRESS_SITE1_APP_PASSWORD=xxxx
|
||||
WORDPRESS_SITE2_ALIAS=myblog
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class SiteConfig(BaseModel):
|
||||
"""
|
||||
Type-safe site configuration.
|
||||
|
||||
Represents configuration for a single site with validation.
|
||||
|
||||
Attributes:
|
||||
site_id: Unique site identifier (e.g., 'site1')
|
||||
plugin_type: Plugin type (e.g., 'wordpress')
|
||||
alias: Optional friendly name
|
||||
config: Site-specific configuration (URL, credentials, etc.)
|
||||
|
||||
Examples:
|
||||
>>> config = SiteConfig(
|
||||
... site_id="site1",
|
||||
... plugin_type="wordpress",
|
||||
... url="https://example.com",
|
||||
... username="admin",
|
||||
... app_password="xxxx"
|
||||
... )
|
||||
"""
|
||||
|
||||
site_id: str = Field(..., description="Unique site identifier")
|
||||
plugin_type: str = Field(..., description="Plugin type (wordpress, gitea, etc)")
|
||||
alias: str | None = Field(None, description="Friendly alias for the site")
|
||||
|
||||
# Common config fields (plugins may require additional fields)
|
||||
url: str | None = Field(None, description="Site URL")
|
||||
username: str | None = Field(None, description="Username for authentication")
|
||||
app_password: str | None = Field(None, description="Application password")
|
||||
container: str | None = Field(None, description="Docker container name (for WP-CLI)")
|
||||
|
||||
# Allow additional fields for plugin-specific configuration
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
@field_validator("alias", mode="before")
|
||||
@classmethod
|
||||
def default_alias(cls, v: str | None, info: ValidationInfo) -> str | None:
|
||||
"""Set alias to site_id if not provided."""
|
||||
return v if v is not None else info.data.get("site_id")
|
||||
|
||||
def get_full_id(self) -> str:
|
||||
"""
|
||||
Get full site identifier.
|
||||
|
||||
Returns:
|
||||
Full ID in format: plugin_type_site_id
|
||||
|
||||
Examples:
|
||||
>>> config.get_full_id()
|
||||
'wordpress_site1'
|
||||
"""
|
||||
return f"{self.plugin_type}_{self.site_id}"
|
||||
|
||||
def get_display_name(self) -> str:
|
||||
"""
|
||||
Get display name for the site.
|
||||
|
||||
Returns:
|
||||
Alias if available, otherwise site_id
|
||||
|
||||
Examples:
|
||||
>>> config.get_display_name()
|
||||
'myblog' # or 'site1' if no alias
|
||||
"""
|
||||
return self.alias or self.site_id
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""
|
||||
Convert to dictionary for plugin consumption.
|
||||
|
||||
Returns:
|
||||
Dictionary with all configuration
|
||||
|
||||
Examples:
|
||||
>>> config_dict = config.to_dict()
|
||||
>>> plugin = WordPressPlugin(config_dict)
|
||||
"""
|
||||
return self.model_dump()
|
||||
|
||||
class SiteManager:
|
||||
"""
|
||||
Manage site configurations with type safety.
|
||||
|
||||
Discovers, validates, and provides access to site configurations.
|
||||
|
||||
Attributes:
|
||||
sites: Dictionary mapping plugin_type to site configurations
|
||||
aliases: Dictionary mapping aliases to (plugin_type, site_id)
|
||||
logger: Logger instance
|
||||
|
||||
Examples:
|
||||
>>> manager = SiteManager()
|
||||
>>> manager.discover_sites(['wordpress', 'gitea'])
|
||||
>>> config = manager.get_site_config('wordpress', 'myblog')
|
||||
>>> sites = manager.list_sites('wordpress')
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize site manager with empty registries."""
|
||||
# Nested dict: plugin_type -> site_id/alias -> SiteConfig
|
||||
self.sites: dict[str, dict[str, SiteConfig]] = {}
|
||||
|
||||
# Map aliases to full_id for quick lookup
|
||||
self.aliases: dict[str, str] = {} # alias -> full_id
|
||||
|
||||
self.logger = logging.getLogger("SiteManager")
|
||||
self.logger.info("SiteManager initialized")
|
||||
|
||||
def discover_sites(self, plugin_types: list[str]) -> int:
|
||||
"""
|
||||
Discover sites from environment variables.
|
||||
|
||||
Scans environment for site configurations and registers them.
|
||||
|
||||
Args:
|
||||
plugin_types: List of plugin types to discover (e.g., ['wordpress'])
|
||||
|
||||
Returns:
|
||||
Number of sites discovered
|
||||
|
||||
Examples:
|
||||
>>> count = manager.discover_sites(['wordpress', 'gitea'])
|
||||
>>> print(f"Discovered {count} sites")
|
||||
"""
|
||||
self.logger.info(f"Starting site discovery for: {', '.join(plugin_types)}")
|
||||
|
||||
total_discovered = 0
|
||||
for plugin_type in plugin_types:
|
||||
count = self._discover_plugin_sites(plugin_type)
|
||||
total_discovered += count
|
||||
|
||||
self.logger.info(
|
||||
f"Discovery complete. Found {total_discovered} sites "
|
||||
f"across {len(plugin_types)} plugin types."
|
||||
)
|
||||
|
||||
return total_discovered
|
||||
|
||||
# Reserved words that should NOT be interpreted as site IDs
|
||||
RESERVED_SITE_WORDS = {
|
||||
"limit",
|
||||
"rate",
|
||||
"config",
|
||||
"debug",
|
||||
"log",
|
||||
"level",
|
||||
"mode",
|
||||
"timeout",
|
||||
"retry",
|
||||
"max",
|
||||
"min",
|
||||
"default",
|
||||
"global",
|
||||
"enabled",
|
||||
"disabled",
|
||||
"host",
|
||||
"port",
|
||||
"path",
|
||||
"key",
|
||||
"secret",
|
||||
"token",
|
||||
"advanced",
|
||||
"basic",
|
||||
"simple",
|
||||
"pro",
|
||||
"premium",
|
||||
"standard",
|
||||
}
|
||||
|
||||
def _discover_plugin_sites(self, plugin_type: str) -> int:
|
||||
"""
|
||||
Discover all sites for a specific plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Type of plugin (e.g., 'wordpress')
|
||||
|
||||
Returns:
|
||||
Number of sites discovered for this plugin type
|
||||
|
||||
Examples:
|
||||
>>> count = manager._discover_plugin_sites('wordpress')
|
||||
"""
|
||||
prefix = plugin_type.upper() + "_"
|
||||
|
||||
# Pattern to match: WORDPRESS_SITE1_URL, WORDPRESS_SITE2_USERNAME, etc.
|
||||
env_pattern = re.compile(f"^{prefix}([A-Z0-9_]+?)_(.+)$")
|
||||
|
||||
# Find all unique site IDs
|
||||
site_ids = set()
|
||||
for env_key in os.environ.keys():
|
||||
match = env_pattern.match(env_key)
|
||||
if match:
|
||||
site_id = match.group(1).lower()
|
||||
# Skip reserved words that are not real site IDs
|
||||
if site_id not in self.RESERVED_SITE_WORDS:
|
||||
site_ids.add(site_id)
|
||||
|
||||
# Load configuration for each site
|
||||
discovered_count = 0
|
||||
for site_id in site_ids:
|
||||
try:
|
||||
config = self._load_site_config(plugin_type, site_id)
|
||||
if config:
|
||||
self.register_site(config)
|
||||
discovered_count += 1
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to load {plugin_type} site '{site_id}': {e}", exc_info=True
|
||||
)
|
||||
|
||||
return discovered_count
|
||||
|
||||
def _load_site_config(self, plugin_type: str, site_id: str) -> SiteConfig | None:
|
||||
"""
|
||||
Load configuration for a site from environment.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type
|
||||
site_id: Site ID
|
||||
|
||||
Returns:
|
||||
SiteConfig if successful, None if incomplete
|
||||
|
||||
Examples:
|
||||
>>> config = manager._load_site_config('wordpress', 'site1')
|
||||
"""
|
||||
prefix = f"{plugin_type.upper()}_{site_id.upper()}_"
|
||||
config_data = {"site_id": site_id, "plugin_type": plugin_type}
|
||||
|
||||
# Collect all config keys for this site
|
||||
for env_key, env_value in os.environ.items():
|
||||
if env_key.startswith(prefix):
|
||||
# Extract config key (everything after prefix)
|
||||
config_key = env_key[len(prefix) :].lower()
|
||||
config_data[config_key] = env_value
|
||||
|
||||
# Must have at least some configuration beyond site_id and plugin_type
|
||||
if len(config_data) <= 2:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Create and validate SiteConfig
|
||||
config = SiteConfig(**config_data)
|
||||
|
||||
self.logger.debug(
|
||||
f"Loaded config for {plugin_type}/{site_id}: " f"{list(config_data.keys())}"
|
||||
)
|
||||
|
||||
return config
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Validation failed for {plugin_type}/{site_id}: {e}", exc_info=True)
|
||||
return None
|
||||
|
||||
def register_site(self, config: SiteConfig) -> None:
|
||||
"""
|
||||
Register a site configuration.
|
||||
|
||||
Args:
|
||||
config: Site configuration to register
|
||||
|
||||
Examples:
|
||||
>>> config = SiteConfig(site_id="site1", plugin_type="wordpress", ...)
|
||||
>>> manager.register_site(config)
|
||||
"""
|
||||
plugin_type = config.plugin_type
|
||||
|
||||
# Initialize plugin_type dict if needed
|
||||
if plugin_type not in self.sites:
|
||||
self.sites[plugin_type] = {}
|
||||
|
||||
# Register by site_id
|
||||
self.sites[plugin_type][config.site_id] = config
|
||||
|
||||
# Register by alias if different from site_id
|
||||
if config.alias and config.alias != config.site_id:
|
||||
self.sites[plugin_type][config.alias] = config
|
||||
# Also register in global aliases map
|
||||
self.aliases[config.alias] = config.get_full_id()
|
||||
|
||||
# Register full_id
|
||||
full_id = config.get_full_id()
|
||||
self.aliases[full_id] = full_id
|
||||
|
||||
# Register site_id
|
||||
self.aliases[config.site_id] = full_id
|
||||
|
||||
self.logger.info(
|
||||
f"Registered site: {full_id} " f"(alias: {config.alias or config.site_id})"
|
||||
)
|
||||
|
||||
def get_site_config(self, plugin_type: str, site: str) -> SiteConfig:
|
||||
"""
|
||||
Get site configuration by ID or alias.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type (e.g., 'wordpress')
|
||||
site: Site ID, alias, or full_id
|
||||
|
||||
Returns:
|
||||
Site configuration
|
||||
|
||||
Raises:
|
||||
ValueError: If site not found
|
||||
|
||||
Examples:
|
||||
>>> config = manager.get_site_config('wordpress', 'myblog')
|
||||
>>> config = manager.get_site_config('wordpress', 'site1')
|
||||
"""
|
||||
if plugin_type not in self.sites:
|
||||
# SECURITY: Don't reveal available plugin types in multi-tenant environment
|
||||
raise ValueError(
|
||||
f"No sites configured for plugin type: {plugin_type}. "
|
||||
f"Please check your environment variables."
|
||||
)
|
||||
|
||||
# Try direct lookup
|
||||
config = self.sites[plugin_type].get(site)
|
||||
if config:
|
||||
return config
|
||||
|
||||
# SECURITY: Don't reveal available sites in multi-tenant environment
|
||||
# Only log available sites count for debugging (not in error message)
|
||||
available_count = len(self.sites[plugin_type])
|
||||
self.logger.debug(
|
||||
f"Site '{site}' not found for {plugin_type}. "
|
||||
f"Total configured sites: {available_count}"
|
||||
)
|
||||
raise ValueError(
|
||||
f"Site '{site}' not configured for {plugin_type}. "
|
||||
f"Please verify the site alias/ID and check environment variables."
|
||||
)
|
||||
|
||||
def list_sites(self, plugin_type: str) -> list[str]:
|
||||
"""
|
||||
List available site IDs and aliases for a plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type
|
||||
|
||||
Returns:
|
||||
List of valid site identifiers (IDs and aliases)
|
||||
|
||||
Examples:
|
||||
>>> sites = manager.list_sites('wordpress')
|
||||
>>> print(sites) # ['site1', 'site2', 'myblog']
|
||||
"""
|
||||
if plugin_type not in self.sites:
|
||||
return []
|
||||
|
||||
# Get unique site identifiers (both IDs and aliases)
|
||||
identifiers = set()
|
||||
for config in self.sites[plugin_type].values():
|
||||
identifiers.add(config.site_id)
|
||||
if config.alias and config.alias != config.site_id:
|
||||
identifiers.add(config.alias)
|
||||
|
||||
# Remove duplicates and sort
|
||||
return sorted(set(identifiers))
|
||||
|
||||
def get_sites_by_type(self, plugin_type: str) -> list[SiteConfig]:
|
||||
"""
|
||||
Get all site configurations for a plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type
|
||||
|
||||
Returns:
|
||||
List of site configurations
|
||||
|
||||
Examples:
|
||||
>>> configs = manager.get_sites_by_type('wordpress')
|
||||
>>> for config in configs:
|
||||
... print(config.get_display_name())
|
||||
"""
|
||||
if plugin_type not in self.sites:
|
||||
return []
|
||||
|
||||
# Return unique configs (since aliases may point to same config)
|
||||
seen_site_ids = set()
|
||||
unique_configs = []
|
||||
|
||||
for config in self.sites[plugin_type].values():
|
||||
if config.site_id not in seen_site_ids:
|
||||
seen_site_ids.add(config.site_id)
|
||||
unique_configs.append(config)
|
||||
|
||||
return unique_configs
|
||||
|
||||
def list_all_sites(self) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List all discovered sites across all plugin types.
|
||||
|
||||
Returns:
|
||||
List of site info dictionaries
|
||||
|
||||
Examples:
|
||||
>>> all_sites = manager.list_all_sites()
|
||||
>>> for site in all_sites:
|
||||
... print(f"{site['full_id']}: {site['alias']}")
|
||||
"""
|
||||
all_sites = []
|
||||
for plugin_type in self.sites:
|
||||
for config in self.get_sites_by_type(plugin_type):
|
||||
all_sites.append(
|
||||
{
|
||||
"plugin_type": config.plugin_type,
|
||||
"site_id": config.site_id,
|
||||
"alias": config.alias,
|
||||
"full_id": config.get_full_id(),
|
||||
}
|
||||
)
|
||||
|
||||
return all_sites
|
||||
|
||||
def get_count(self) -> int:
|
||||
"""
|
||||
Get total number of registered sites.
|
||||
|
||||
Returns:
|
||||
Total site count
|
||||
|
||||
Examples:
|
||||
>>> count = manager.get_count()
|
||||
"""
|
||||
total = 0
|
||||
for plugin_type in self.sites:
|
||||
total += len(self.get_sites_by_type(plugin_type))
|
||||
return total
|
||||
|
||||
def get_count_by_type(self) -> dict[str, int]:
|
||||
"""
|
||||
Get site counts grouped by plugin type.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping plugin type to site count
|
||||
|
||||
Examples:
|
||||
>>> counts = manager.get_count_by_type()
|
||||
>>> print(counts) # {'wordpress': 4, 'gitea': 2}
|
||||
"""
|
||||
return {plugin_type: len(self.get_sites_by_type(plugin_type)) for plugin_type in self.sites}
|
||||
|
||||
def get_effective_path_suffix(self, full_id: str) -> str:
|
||||
"""
|
||||
Get the effective path suffix for a site's endpoint.
|
||||
|
||||
Uses alias if available, otherwise returns full_id.
|
||||
|
||||
Args:
|
||||
full_id: The full site ID (e.g., 'wordpress_site1')
|
||||
|
||||
Returns:
|
||||
Path suffix to use in endpoint URL (alias or full_id)
|
||||
|
||||
Examples:
|
||||
>>> suffix = manager.get_effective_path_suffix('wordpress_site1')
|
||||
>>> print(suffix) # 'myblog' or 'wordpress_site1'
|
||||
"""
|
||||
# Parse full_id to get plugin_type and site_id
|
||||
parts = full_id.split("_", 1)
|
||||
if len(parts) != 2:
|
||||
return full_id
|
||||
|
||||
plugin_type, site_id = parts
|
||||
|
||||
# Try to get config
|
||||
try:
|
||||
config = self.get_site_config(plugin_type, site_id)
|
||||
# If alias exists and is different from site_id, use alias
|
||||
if config.alias and config.alias != config.site_id:
|
||||
return config.alias
|
||||
return full_id
|
||||
except ValueError:
|
||||
return full_id
|
||||
|
||||
def get_alias_conflicts(self) -> dict[str, list[str]]:
|
||||
"""
|
||||
Get all alias conflicts.
|
||||
|
||||
Note: SiteManager doesn't track conflicts like SiteRegistry.
|
||||
This is a stub for compatibility.
|
||||
|
||||
Returns:
|
||||
Empty dict (no conflict tracking in SiteManager)
|
||||
"""
|
||||
return {}
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""String representation of site manager."""
|
||||
counts = self.get_count_by_type()
|
||||
counts_str = ", ".join(f"{k}: {v}" for k, v in counts.items())
|
||||
return f"SiteManager(total={self.get_count()}, {counts_str})"
|
||||
|
||||
# Singleton instance
|
||||
_site_manager: SiteManager | None = None
|
||||
|
||||
def get_site_manager() -> SiteManager:
|
||||
"""
|
||||
Get the singleton site manager instance.
|
||||
|
||||
Returns:
|
||||
Global SiteManager instance
|
||||
|
||||
Examples:
|
||||
>>> manager = get_site_manager()
|
||||
>>> manager.discover_sites(['wordpress'])
|
||||
"""
|
||||
global _site_manager
|
||||
if _site_manager is None:
|
||||
_site_manager = SiteManager()
|
||||
return _site_manager
|
||||
371
core/site_registry.py
Normal file
371
core/site_registry.py
Normal file
@@ -0,0 +1,371 @@
|
||||
"""
|
||||
Site Registry
|
||||
|
||||
Manages site configurations discovered from environment variables.
|
||||
Supports multi-site setups with aliases and dynamic discovery.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class SiteInfo:
|
||||
"""Information about a single site."""
|
||||
|
||||
def __init__(
|
||||
self, plugin_type: str, site_id: str, config: dict[str, Any], alias: str | None = None
|
||||
):
|
||||
"""
|
||||
Initialize site information.
|
||||
|
||||
Args:
|
||||
plugin_type: Type of plugin (e.g., 'wordpress')
|
||||
site_id: Site identifier (e.g., 'site1')
|
||||
config: Site configuration from environment
|
||||
alias: Optional friendly alias for the site
|
||||
"""
|
||||
self.plugin_type = plugin_type
|
||||
self.site_id = site_id
|
||||
self.config = config
|
||||
self.alias = alias or site_id
|
||||
|
||||
def get_full_id(self) -> str:
|
||||
"""Get full site identifier: plugin_type_site_id"""
|
||||
return f"{self.plugin_type}_{self.site_id}"
|
||||
|
||||
def get_display_name(self) -> str:
|
||||
"""Get display name (alias if available, otherwise site_id)"""
|
||||
return self.alias
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"plugin_type": self.plugin_type,
|
||||
"site_id": self.site_id,
|
||||
"alias": self.alias,
|
||||
"full_id": self.get_full_id(),
|
||||
"config_keys": list(self.config.keys()),
|
||||
}
|
||||
|
||||
class SiteRegistry:
|
||||
"""
|
||||
Registry for managing site configurations across plugin types.
|
||||
|
||||
Discovers sites from environment variables:
|
||||
- {PLUGIN_TYPE}_{SITE_ID}_{CONFIG_KEY}
|
||||
- {PLUGIN_TYPE}_{SITE_ID}_ALIAS (optional)
|
||||
|
||||
Example:
|
||||
WORDPRESS_SITE1_URL=https://example.com
|
||||
WORDPRESS_SITE1_USERNAME=admin
|
||||
WORDPRESS_SITE1_APP_PASSWORD=xxxx
|
||||
WORDPRESS_SITE2_URL=https://myblog.com
|
||||
WORDPRESS_SITE2_ALIAS=myblog
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize site registry."""
|
||||
self.sites: dict[str, SiteInfo] = {} # full_id -> SiteInfo
|
||||
self.aliases: dict[str, str] = {} # alias -> full_id
|
||||
self.alias_conflicts: dict[str, list[str]] = {} # alias -> [full_ids that wanted it]
|
||||
self.logger = logging.getLogger("SiteRegistry")
|
||||
|
||||
def discover_sites(self, plugin_types: list[str]) -> None:
|
||||
"""
|
||||
Discover sites from environment variables.
|
||||
|
||||
Args:
|
||||
plugin_types: List of plugin types to discover (e.g., ['wordpress'])
|
||||
"""
|
||||
self.logger.info("Starting site discovery...")
|
||||
|
||||
for plugin_type in plugin_types:
|
||||
self._discover_plugin_sites(plugin_type)
|
||||
|
||||
self.logger.info(
|
||||
f"Discovery complete. Found {len(self.sites)} sites "
|
||||
f"with {len(self.aliases)} aliases."
|
||||
)
|
||||
|
||||
# Log alias conflicts if any
|
||||
if self.alias_conflicts:
|
||||
self.logger.warning("=" * 50)
|
||||
self.logger.warning("DUPLICATE ALIAS CONFLICTS DETECTED:")
|
||||
for alias, full_ids in self.alias_conflicts.items():
|
||||
winner = self.aliases.get(alias)
|
||||
losers = [fid for fid in full_ids if fid != winner]
|
||||
self.logger.warning(
|
||||
f" Alias '{alias}': {winner} (winner), {losers} (using full_id)"
|
||||
)
|
||||
self.logger.warning("=" * 50)
|
||||
|
||||
# Reserved words that should NOT be interpreted as site IDs
|
||||
RESERVED_SITE_WORDS = {
|
||||
"limit",
|
||||
"rate",
|
||||
"config",
|
||||
"debug",
|
||||
"log",
|
||||
"level",
|
||||
"mode",
|
||||
"timeout",
|
||||
"retry",
|
||||
"max",
|
||||
"min",
|
||||
"default",
|
||||
"global",
|
||||
"enabled",
|
||||
"disabled",
|
||||
"host",
|
||||
"port",
|
||||
"path",
|
||||
"key",
|
||||
"secret",
|
||||
"token",
|
||||
"advanced",
|
||||
"basic",
|
||||
"simple",
|
||||
"pro",
|
||||
"premium",
|
||||
"standard",
|
||||
}
|
||||
|
||||
def _discover_plugin_sites(self, plugin_type: str) -> None:
|
||||
"""
|
||||
Discover all sites for a specific plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Type of plugin (e.g., 'wordpress')
|
||||
"""
|
||||
prefix = plugin_type.upper() + "_"
|
||||
|
||||
# Find all site IDs for this plugin type
|
||||
site_ids = set()
|
||||
env_pattern = re.compile(f"^{prefix}([A-Z0-9_]+?)_(.+)$")
|
||||
|
||||
for env_key in os.environ.keys():
|
||||
match = env_pattern.match(env_key)
|
||||
if match:
|
||||
site_id = match.group(1).lower()
|
||||
# Skip reserved words that are not real site IDs
|
||||
if site_id not in self.RESERVED_SITE_WORDS:
|
||||
site_ids.add(site_id)
|
||||
|
||||
# Create SiteInfo for each discovered site
|
||||
for site_id in site_ids:
|
||||
try:
|
||||
config = self._load_site_config(plugin_type, site_id)
|
||||
if config:
|
||||
alias = config.pop("alias", None) # Extract alias if present
|
||||
site_info = SiteInfo(plugin_type, site_id, config, alias)
|
||||
|
||||
full_id = site_info.get_full_id()
|
||||
self.sites[full_id] = site_info
|
||||
|
||||
# Register alias with duplicate detection
|
||||
if alias:
|
||||
self._register_alias_safe(alias, full_id)
|
||||
# Also register with prefix
|
||||
prefixed_alias = f"{plugin_type}_{alias}"
|
||||
self._register_alias_safe(prefixed_alias, full_id)
|
||||
|
||||
# Always register site_id as alias too (with safe check)
|
||||
self._register_alias_safe(site_id, full_id)
|
||||
self.aliases[full_id] = (
|
||||
full_id # full_id can reference itself (no conflict possible)
|
||||
)
|
||||
|
||||
# Log with alias status
|
||||
effective_alias = self._get_effective_alias(alias, full_id)
|
||||
self.logger.info(f"Discovered site: {full_id} " f"(path: {effective_alias})")
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to load {plugin_type} site '{site_id}': {e}", exc_info=True
|
||||
)
|
||||
|
||||
def _register_alias_safe(self, alias: str, full_id: str) -> bool:
|
||||
"""
|
||||
Register an alias with duplicate detection.
|
||||
|
||||
Args:
|
||||
alias: The alias to register
|
||||
full_id: The full_id to map the alias to
|
||||
|
||||
Returns:
|
||||
True if alias was registered, False if it was a duplicate
|
||||
"""
|
||||
if alias in self.aliases:
|
||||
existing_full_id = self.aliases[alias]
|
||||
if existing_full_id != full_id:
|
||||
# Duplicate alias detected
|
||||
if alias not in self.alias_conflicts:
|
||||
self.alias_conflicts[alias] = [existing_full_id]
|
||||
self.alias_conflicts[alias].append(full_id)
|
||||
self.logger.warning(
|
||||
f"Duplicate alias '{alias}': {full_id} conflicts with {existing_full_id}. "
|
||||
f"{full_id} will use full_id for endpoint path."
|
||||
)
|
||||
return False
|
||||
else:
|
||||
self.aliases[alias] = full_id
|
||||
return True
|
||||
|
||||
def _get_effective_alias(self, alias: str | None, full_id: str) -> str:
|
||||
"""
|
||||
Get the effective path suffix for a site.
|
||||
|
||||
If the alias was taken by another site, returns full_id.
|
||||
Otherwise returns the alias (or full_id if no alias).
|
||||
|
||||
Args:
|
||||
alias: The desired alias
|
||||
full_id: The full site ID
|
||||
|
||||
Returns:
|
||||
The effective path suffix
|
||||
"""
|
||||
if alias:
|
||||
# Check if this site owns the alias
|
||||
if self.aliases.get(alias) == full_id:
|
||||
return alias
|
||||
else:
|
||||
# Alias was taken, use full_id
|
||||
return full_id
|
||||
return full_id
|
||||
|
||||
def get_alias_conflicts(self) -> dict[str, list[str]]:
|
||||
"""
|
||||
Get all alias conflicts.
|
||||
|
||||
Returns:
|
||||
Dict mapping conflicted aliases to list of full_ids that wanted them
|
||||
"""
|
||||
return self.alias_conflicts.copy()
|
||||
|
||||
def get_effective_path_suffix(self, full_id: str) -> str:
|
||||
"""
|
||||
Get the effective path suffix for a site's endpoint.
|
||||
|
||||
Uses alias if available and not conflicted, otherwise full_id.
|
||||
|
||||
Args:
|
||||
full_id: The full site ID
|
||||
|
||||
Returns:
|
||||
Path suffix to use in endpoint URL
|
||||
"""
|
||||
site_info = self.sites.get(full_id)
|
||||
if not site_info:
|
||||
return full_id
|
||||
|
||||
alias = site_info.alias
|
||||
return self._get_effective_alias(alias, full_id)
|
||||
|
||||
def _load_site_config(self, plugin_type: str, site_id: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Load configuration for a site from environment.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type
|
||||
site_id: Site ID
|
||||
|
||||
Returns:
|
||||
Dict with configuration or None if incomplete
|
||||
"""
|
||||
prefix = f"{plugin_type.upper()}_{site_id.upper()}_"
|
||||
config = {}
|
||||
|
||||
# Collect all config keys for this site
|
||||
for env_key, env_value in os.environ.items():
|
||||
if env_key.startswith(prefix):
|
||||
# Extract config key (everything after prefix)
|
||||
config_key = env_key[len(prefix) :].lower()
|
||||
config[config_key] = env_value
|
||||
|
||||
if not config:
|
||||
return None
|
||||
|
||||
self.logger.debug(f"Loaded config for {plugin_type}/{site_id}: {list(config.keys())}")
|
||||
return config
|
||||
|
||||
def get_site(self, plugin_type: str, site_identifier: str) -> SiteInfo | None:
|
||||
"""
|
||||
Get site information by ID or alias.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type (e.g., 'wordpress')
|
||||
site_identifier: Site ID, alias, or full_id
|
||||
|
||||
Returns:
|
||||
SiteInfo if found, None otherwise
|
||||
"""
|
||||
# Try direct lookup first
|
||||
full_id = f"{plugin_type}_{site_identifier}"
|
||||
if full_id in self.sites:
|
||||
return self.sites[full_id]
|
||||
|
||||
# Try alias lookup
|
||||
if site_identifier in self.aliases:
|
||||
resolved_full_id = self.aliases[site_identifier]
|
||||
return self.sites.get(resolved_full_id)
|
||||
|
||||
# Try prefixed alias
|
||||
prefixed = f"{plugin_type}_{site_identifier}"
|
||||
if prefixed in self.aliases:
|
||||
resolved_full_id = self.aliases[prefixed]
|
||||
return self.sites.get(resolved_full_id)
|
||||
|
||||
return None
|
||||
|
||||
def get_sites_by_type(self, plugin_type: str) -> list[SiteInfo]:
|
||||
"""
|
||||
Get all sites of a specific plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type to filter by
|
||||
|
||||
Returns:
|
||||
List of SiteInfo objects
|
||||
"""
|
||||
return [
|
||||
site_info for site_info in self.sites.values() if site_info.plugin_type == plugin_type
|
||||
]
|
||||
|
||||
def list_all_sites(self) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List all discovered sites.
|
||||
|
||||
Returns:
|
||||
List of site info dictionaries
|
||||
"""
|
||||
return [site_info.to_dict() for site_info in self.sites.values()]
|
||||
|
||||
def get_site_options(self, plugin_type: str) -> list[str]:
|
||||
"""
|
||||
Get available site options for a plugin type (for schema enum).
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type
|
||||
|
||||
Returns:
|
||||
List of valid site identifiers (IDs and aliases)
|
||||
"""
|
||||
options = set()
|
||||
for site_info in self.get_sites_by_type(plugin_type):
|
||||
options.add(site_info.site_id)
|
||||
if site_info.alias and site_info.alias != site_info.site_id:
|
||||
options.add(site_info.alias)
|
||||
return sorted(options)
|
||||
|
||||
# Global site registry instance
|
||||
_site_registry: SiteRegistry | None = None
|
||||
|
||||
def get_site_registry() -> SiteRegistry:
|
||||
"""Get the global site registry instance."""
|
||||
global _site_registry
|
||||
if _site_registry is None:
|
||||
_site_registry = SiteRegistry()
|
||||
return _site_registry
|
||||
511
core/tool_generator.py
Normal file
511
core/tool_generator.py
Normal file
@@ -0,0 +1,511 @@
|
||||
"""
|
||||
Tool Generator - Direct tool generation without per-site wrapper
|
||||
|
||||
Generates MCP tools directly from plugin specifications.
|
||||
Part of Option B clean architecture refactoring.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from core.tool_registry import ToolDefinition
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Plugin type fallback mapping - used when a plugin has no sites configured
|
||||
# WooCommerce can fallback to WordPress sites (same URL, credentials)
|
||||
# NOTE: Using fallback is NOT recommended in production. Always configure
|
||||
# explicit WOOCOMMERCE_SITE*_... environment variables for stability.
|
||||
PLUGIN_SITE_FALLBACK = {
|
||||
"woocommerce": "wordpress", # WooCommerce can use WordPress site configs
|
||||
# Add more fallbacks as needed
|
||||
}
|
||||
|
||||
def get_site_plugin_type_with_fallback(plugin_type: str, site_manager) -> str:
|
||||
"""
|
||||
Get the site configuration plugin type for a given plugin.
|
||||
|
||||
Checks if the plugin has its own sites configured first.
|
||||
If not, falls back to a related plugin type (e.g., woocommerce -> wordpress).
|
||||
|
||||
WARNING: Using fallback is not recommended for production use.
|
||||
Always configure explicit environment variables for each plugin type
|
||||
to avoid issues with alias mismatches and credential problems.
|
||||
|
||||
Args:
|
||||
plugin_type: The plugin type
|
||||
site_manager: SiteManager instance to check for configured sites
|
||||
|
||||
Returns:
|
||||
The plugin type to use for site configuration lookup
|
||||
"""
|
||||
# First check if the plugin has its own sites
|
||||
if plugin_type in site_manager.sites and site_manager.sites[plugin_type]:
|
||||
return plugin_type
|
||||
|
||||
# Fallback to related plugin type if available
|
||||
fallback_type = PLUGIN_SITE_FALLBACK.get(plugin_type)
|
||||
if fallback_type and fallback_type in site_manager.sites:
|
||||
# Log a warning - fallback usage is not recommended
|
||||
logger.warning(
|
||||
f"FALLBACK: Using {fallback_type} site config for {plugin_type}. "
|
||||
f"This is NOT recommended for production. "
|
||||
f"Please configure explicit {plugin_type.upper()}_SITE*_... environment variables "
|
||||
f"to avoid potential issues with alias mismatches and credentials."
|
||||
)
|
||||
return fallback_type
|
||||
|
||||
# Return original type (may have no sites)
|
||||
return plugin_type
|
||||
|
||||
class ToolGenerator:
|
||||
"""
|
||||
Generate tools directly from plugin classes.
|
||||
|
||||
No longer wraps per-site tools - generates tools directly
|
||||
from plugin specifications with site parameter routing.
|
||||
|
||||
Attributes:
|
||||
site_manager: Manages site configurations
|
||||
logger: Logger instance
|
||||
|
||||
Examples:
|
||||
>>> generator = ToolGenerator(site_manager)
|
||||
>>> tools = generator.generate_tools(WordPressPlugin, "wordpress")
|
||||
>>> print(f"Generated {len(tools)} tools")
|
||||
"""
|
||||
|
||||
def __init__(self, site_manager):
|
||||
"""
|
||||
Initialize tool generator.
|
||||
|
||||
Args:
|
||||
site_manager: SiteManager instance for site configuration
|
||||
"""
|
||||
self.site_manager = site_manager
|
||||
self.logger = logging.getLogger("ToolGenerator")
|
||||
|
||||
def generate_tools(self, plugin_class: type, plugin_type: str) -> list[ToolDefinition]:
|
||||
"""
|
||||
Generate tools directly from plugin class.
|
||||
|
||||
Args:
|
||||
plugin_class: Plugin class (not instance)
|
||||
plugin_type: Plugin type name (e.g., 'wordpress')
|
||||
|
||||
Returns:
|
||||
List of tool definitions
|
||||
|
||||
Raises:
|
||||
ValueError: If plugin_class doesn't have get_tool_specifications
|
||||
|
||||
Examples:
|
||||
>>> tools = generator.generate_tools(WordPressPlugin, "wordpress")
|
||||
"""
|
||||
# Verify plugin class has required method
|
||||
if not hasattr(plugin_class, "get_tool_specifications"):
|
||||
raise ValueError(
|
||||
f"Plugin class {plugin_class.__name__} must implement "
|
||||
"get_tool_specifications() static method"
|
||||
)
|
||||
|
||||
# Get tool specifications from plugin
|
||||
try:
|
||||
tool_specs = plugin_class.get_tool_specifications()
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to get tool specifications from {plugin_class.__name__}: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
return []
|
||||
|
||||
self.logger.info(
|
||||
f"Generating tools for {plugin_type} " f"from {len(tool_specs)} specifications"
|
||||
)
|
||||
|
||||
tools = []
|
||||
for spec in tool_specs:
|
||||
try:
|
||||
tool = self._create_tool_from_spec(plugin_class, plugin_type, spec)
|
||||
if tool:
|
||||
tools.append(tool)
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to create tool from spec {spec.get('name', 'unknown')}: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
self.logger.info(f"Generated {len(tools)} tools for {plugin_type}")
|
||||
return tools
|
||||
|
||||
def _create_tool_from_spec(
|
||||
self, plugin_class: type, plugin_type: str, spec: dict[str, Any]
|
||||
) -> ToolDefinition:
|
||||
"""
|
||||
Create a tool definition from a specification.
|
||||
|
||||
Args:
|
||||
plugin_class: Plugin class
|
||||
plugin_type: Plugin type name
|
||||
spec: Tool specification dictionary
|
||||
|
||||
Returns:
|
||||
Tool definition
|
||||
|
||||
Raises:
|
||||
ValueError: If spec is invalid
|
||||
|
||||
Tool spec format:
|
||||
{
|
||||
'name': 'list_posts',
|
||||
'method_name': 'list_posts',
|
||||
'description': 'List WordPress posts',
|
||||
'schema': {...}, # Input schema (without site param)
|
||||
'scope': 'read' # Optional, defaults to 'read'
|
||||
}
|
||||
"""
|
||||
# Validate required fields
|
||||
required_fields = ["name", "method_name", "description", "schema"]
|
||||
for field in required_fields:
|
||||
if field not in spec:
|
||||
raise ValueError(f"Tool spec missing required field: {field}")
|
||||
|
||||
# Extract spec fields
|
||||
action_name = spec["name"]
|
||||
method_name = spec["method_name"]
|
||||
description = spec["description"]
|
||||
schema = spec["schema"]
|
||||
scope = spec.get("scope", "read")
|
||||
|
||||
# Create full tool name
|
||||
tool_name = f"{plugin_type}_{action_name}"
|
||||
|
||||
# Add site parameter to schema
|
||||
enhanced_schema = self._add_site_parameter(schema, plugin_type)
|
||||
|
||||
# Add [UNIFIED] prefix to description if not present
|
||||
if not description.startswith("[UNIFIED]"):
|
||||
description = f"[UNIFIED] {description}"
|
||||
|
||||
# Create handler with site routing
|
||||
handler = self._create_handler(plugin_class, plugin_type, method_name)
|
||||
|
||||
return ToolDefinition(
|
||||
name=tool_name,
|
||||
description=description,
|
||||
input_schema=enhanced_schema,
|
||||
handler=handler,
|
||||
required_scope=scope,
|
||||
plugin_type=plugin_type,
|
||||
)
|
||||
|
||||
def _add_site_parameter(
|
||||
self, original_schema: dict[str, Any], plugin_type: str
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Add 'site' parameter to input schema.
|
||||
|
||||
Args:
|
||||
original_schema: Original input schema
|
||||
plugin_type: Plugin type for site options
|
||||
|
||||
Returns:
|
||||
Schema with site parameter added
|
||||
|
||||
Examples:
|
||||
>>> schema = {'type': 'object', 'properties': {'post_id': {...}}}
|
||||
>>> enhanced = generator._add_site_parameter(schema, 'wordpress')
|
||||
>>> assert 'site' in enhanced['properties']
|
||||
"""
|
||||
# Deep copy to avoid modifying original
|
||||
schema = copy.deepcopy(original_schema)
|
||||
|
||||
# Ensure schema has required structure
|
||||
if "properties" not in schema:
|
||||
schema["properties"] = {}
|
||||
if "required" not in schema:
|
||||
schema["required"] = []
|
||||
|
||||
# Get available sites for this plugin type
|
||||
# Use fallback if no sites configured (e.g., woocommerce -> wordpress)
|
||||
site_plugin_type = get_site_plugin_type_with_fallback(plugin_type, self.site_manager)
|
||||
site_options = self.site_manager.list_sites(site_plugin_type)
|
||||
|
||||
if not site_options:
|
||||
# No sites configured - add site param anyway for future use
|
||||
site_options = []
|
||||
|
||||
# Phase K.2.6: For single-site MCPs, make site parameter optional
|
||||
is_single_site = len(site_options) == 1
|
||||
|
||||
if is_single_site:
|
||||
# Single site - parameter is optional with helpful description
|
||||
single_site = site_options[0]
|
||||
schema["properties"] = {
|
||||
"site": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
f"🔗 SINGLE SITE: Connected to '{single_site}'. "
|
||||
f"This parameter is OPTIONAL - you can omit it or use any value."
|
||||
),
|
||||
"default": single_site,
|
||||
},
|
||||
**schema["properties"],
|
||||
}
|
||||
# Don't add 'site' to required for single-site MCPs
|
||||
else:
|
||||
# Multiple sites or no sites - parameter is required
|
||||
schema["properties"] = {
|
||||
"site": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
f"Site ID or alias. "
|
||||
f"Available options: {', '.join(site_options) if site_options else 'None configured'}. "
|
||||
f"Use list_sites() to see all configured sites."
|
||||
),
|
||||
**({"enum": site_options} if site_options else {}),
|
||||
},
|
||||
**schema["properties"],
|
||||
}
|
||||
# Make 'site' required for multi-site MCPs
|
||||
if "site" not in schema["required"]:
|
||||
schema["required"].insert(0, "site")
|
||||
|
||||
return schema
|
||||
|
||||
def _create_handler(self, plugin_class: type, plugin_type: str, method_name: str) -> Callable:
|
||||
"""
|
||||
Create async handler with site routing.
|
||||
|
||||
The handler:
|
||||
1. Extracts site from parameters
|
||||
2. Gets site configuration
|
||||
3. Creates plugin instance for this request
|
||||
4. Calls the specified method
|
||||
5. Returns result
|
||||
|
||||
Args:
|
||||
plugin_class: Plugin class to instantiate
|
||||
plugin_type: Plugin type name
|
||||
method_name: Method name to call on plugin instance
|
||||
|
||||
Returns:
|
||||
Async handler function
|
||||
|
||||
Examples:
|
||||
>>> handler = generator._create_handler(
|
||||
... WordPressPlugin, "wordpress", "list_posts"
|
||||
... )
|
||||
>>> result = await handler(site="site1", per_page=10)
|
||||
"""
|
||||
|
||||
async def unified_handler(site: str = None, **kwargs):
|
||||
"""
|
||||
Unified handler that routes to the correct site plugin.
|
||||
|
||||
Args:
|
||||
site: Site ID or alias (optional for single-site MCPs)
|
||||
**kwargs: Other parameters for the tool
|
||||
|
||||
Returns:
|
||||
Result from plugin method
|
||||
|
||||
Raises:
|
||||
ValueError: If site not found or access denied
|
||||
"""
|
||||
try:
|
||||
# Get site configuration
|
||||
# Use fallback if no sites configured (e.g., woocommerce -> wordpress)
|
||||
site_plugin_type = get_site_plugin_type_with_fallback(
|
||||
plugin_type, self.site_manager
|
||||
)
|
||||
|
||||
# Phase K.2.6: Auto-detect site for single-site MCPs
|
||||
if not site:
|
||||
available_sites = self.site_manager.list_sites(site_plugin_type)
|
||||
if len(available_sites) == 1:
|
||||
site = available_sites[0]
|
||||
elif len(available_sites) == 0:
|
||||
return "Error: No sites configured. Please check environment variables."
|
||||
else:
|
||||
return (
|
||||
f"Error: Multiple sites available ({', '.join(available_sites)}). "
|
||||
f"Please specify the 'site' parameter."
|
||||
)
|
||||
|
||||
site_config = self.site_manager.get_site_config(site_plugin_type, site)
|
||||
|
||||
# SECURITY: Check if API key has access to this project
|
||||
from core.context import get_api_key_context
|
||||
|
||||
api_key_info = get_api_key_context()
|
||||
|
||||
if api_key_info and not api_key_info.get("is_global"):
|
||||
# Per-project key - must match the project
|
||||
allowed_project = api_key_info.get("project_id")
|
||||
|
||||
# Resolve the current project - always use site_id for consistency
|
||||
# Use site_plugin_type for consistent project naming
|
||||
current_project = f"{site_plugin_type}_{site_config.site_id}"
|
||||
|
||||
# Resolve allowed_project to normalize alias vs site_id
|
||||
# API key might have been created with alias (wordpress_myblog)
|
||||
# or site_id (wordpress_site1)
|
||||
allowed_project_normalized = allowed_project
|
||||
if allowed_project and "_" in allowed_project:
|
||||
# Extract plugin type and site identifier from allowed_project
|
||||
allowed_parts = allowed_project.split("_", 1)
|
||||
if len(allowed_parts) == 2:
|
||||
allowed_plugin_type, allowed_site_identifier = allowed_parts
|
||||
# Try to resolve the site identifier to site_id
|
||||
try:
|
||||
allowed_site_config = self.site_manager.get_site_config(
|
||||
allowed_plugin_type, allowed_site_identifier
|
||||
)
|
||||
# Normalize to plugin_type_site_id format
|
||||
allowed_project_normalized = (
|
||||
f"{allowed_plugin_type}_{allowed_site_config.site_id}"
|
||||
)
|
||||
except ValueError:
|
||||
# Site not found, keep original for error message
|
||||
pass
|
||||
|
||||
if allowed_project_normalized != current_project:
|
||||
logger.warning(
|
||||
f"Access denied: API key for project '{allowed_project}' "
|
||||
f"attempted to access '{current_project}'"
|
||||
)
|
||||
return (
|
||||
f"Error: Access denied. This API key is restricted to project '{allowed_project}'. "
|
||||
f"Use a global API key or create a key for '{current_project}'."
|
||||
)
|
||||
|
||||
# Create plugin instance for this request
|
||||
# Convert SiteConfig Pydantic model to dict
|
||||
# Use model_dump() for Pydantic V2, fallback to dict() for V1
|
||||
if hasattr(site_config, "model_dump"):
|
||||
config_dict = site_config.model_dump()
|
||||
elif hasattr(site_config, "dict"):
|
||||
config_dict = site_config.dict()
|
||||
else:
|
||||
config_dict = site_config
|
||||
|
||||
plugin_instance = plugin_class(config_dict)
|
||||
|
||||
# Get the method from plugin instance
|
||||
if not hasattr(plugin_instance, method_name):
|
||||
return (
|
||||
f"Error: Method '{method_name}' not found in "
|
||||
f"{plugin_class.__name__}. This is a plugin implementation error."
|
||||
)
|
||||
|
||||
method = getattr(plugin_instance, method_name)
|
||||
|
||||
# Phase K.2.1: Enhanced parameter processing
|
||||
# 1. Parse JSON strings to objects (for billing, shipping, line_items, etc.)
|
||||
# 2. Filter out None and empty values
|
||||
import json as json_module
|
||||
|
||||
def process_value(value):
|
||||
"""Process parameter value - parse JSON strings if needed."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
# Skip empty strings
|
||||
if value.strip() == "":
|
||||
return None
|
||||
# Try to parse JSON strings
|
||||
stripped = value.strip()
|
||||
if (stripped.startswith("{") and stripped.endswith("}")) or (
|
||||
stripped.startswith("[") and stripped.endswith("]")
|
||||
):
|
||||
try:
|
||||
return json_module.loads(value)
|
||||
except json_module.JSONDecodeError:
|
||||
# Not valid JSON, return original string
|
||||
return value
|
||||
return value
|
||||
|
||||
filtered_kwargs = {}
|
||||
for key, value in kwargs.items():
|
||||
processed = process_value(value)
|
||||
if processed is not None:
|
||||
filtered_kwargs[key] = processed
|
||||
|
||||
# Call the method
|
||||
result = await method(**filtered_kwargs)
|
||||
return result
|
||||
|
||||
except ValueError as e:
|
||||
# Site not found or validation error
|
||||
logger.warning(f"Validation error in {plugin_type}_{method_name}: {e}")
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
except Exception as e:
|
||||
# Import custom exceptions for better error handling
|
||||
from plugins.wordpress.client import AuthenticationError, ConfigurationError
|
||||
|
||||
error_type = type(e).__name__
|
||||
|
||||
if isinstance(e, ConfigurationError):
|
||||
# Configuration error - likely missing env vars
|
||||
logger.error(f"Configuration error in {plugin_type}_{method_name}: {e}")
|
||||
return (
|
||||
f"Configuration Error: {str(e)}\n\n"
|
||||
f"Hint: For {plugin_type}, ensure these environment variables are set:\n"
|
||||
f" - {plugin_type.upper()}_SITE*_URL\n"
|
||||
f" - {plugin_type.upper()}_SITE*_USERNAME\n"
|
||||
f" - {plugin_type.upper()}_SITE*_APP_PASSWORD"
|
||||
)
|
||||
|
||||
elif isinstance(e, AuthenticationError):
|
||||
# Authentication error - 401/403
|
||||
logger.warning(f"Authentication error in {plugin_type}_{method_name}: {e}")
|
||||
return f"Authentication Error: {str(e)}"
|
||||
|
||||
else:
|
||||
# Unexpected error
|
||||
logger.error(
|
||||
f"Error in unified handler for {plugin_type}_{method_name}: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
return f"Error ({error_type}): {str(e)}"
|
||||
|
||||
# Set function name for better debugging
|
||||
unified_handler.__name__ = f"{plugin_type}_{method_name}_handler"
|
||||
|
||||
return unified_handler
|
||||
|
||||
def generate_all_tools(self, plugin_classes: dict[str, type]) -> list[ToolDefinition]:
|
||||
"""
|
||||
Generate tools for all plugin classes.
|
||||
|
||||
Args:
|
||||
plugin_classes: Dict mapping plugin_type to plugin class
|
||||
|
||||
Returns:
|
||||
List of all tool definitions
|
||||
|
||||
Examples:
|
||||
>>> plugins = {
|
||||
... 'wordpress': WordPressPlugin,
|
||||
... 'gitea': GiteaPlugin
|
||||
... }
|
||||
>>> all_tools = generator.generate_all_tools(plugins)
|
||||
"""
|
||||
all_tools = []
|
||||
|
||||
for plugin_type, plugin_class in plugin_classes.items():
|
||||
try:
|
||||
tools = self.generate_tools(plugin_class, plugin_type)
|
||||
all_tools.extend(tools)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to generate tools for {plugin_type}: {e}", exc_info=True)
|
||||
|
||||
self.logger.info(
|
||||
f"Generated {len(all_tools)} total tools " f"across {len(plugin_classes)} plugin types"
|
||||
)
|
||||
|
||||
return all_tools
|
||||
228
core/tool_registry.py
Normal file
228
core/tool_registry.py
Normal file
@@ -0,0 +1,228 @@
|
||||
"""
|
||||
Tool Registry - Central tool management for MCP
|
||||
|
||||
Manages tool definitions with type safety and validation.
|
||||
Part of Option B clean architecture refactoring.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ToolDefinition(BaseModel):
|
||||
"""
|
||||
Type-safe tool definition.
|
||||
|
||||
Represents a single MCP tool with all necessary metadata.
|
||||
|
||||
Attributes:
|
||||
name: Unique tool identifier (e.g., "wordpress_get_post")
|
||||
description: Human-readable tool description
|
||||
input_schema: JSON Schema for tool parameters
|
||||
handler: Async function that executes the tool
|
||||
required_scope: Required API key scope ("read", "write", "admin")
|
||||
plugin_type: Plugin type this tool belongs to (e.g., "wordpress")
|
||||
"""
|
||||
|
||||
name: str = Field(..., description="Unique tool identifier")
|
||||
description: str = Field(..., description="Tool description")
|
||||
input_schema: dict[str, Any] = Field(
|
||||
default_factory=lambda: {"type": "object", "properties": {}},
|
||||
description="JSON Schema for parameters",
|
||||
)
|
||||
handler: Callable = Field(..., description="Async handler function")
|
||||
required_scope: str = Field(
|
||||
default="read", description="Required API key scope (read/write/admin)"
|
||||
)
|
||||
plugin_type: str = Field(..., description="Plugin type (wordpress, gitea, etc)")
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True) # Allow Callable type
|
||||
|
||||
class ToolRegistry:
|
||||
"""
|
||||
Central registry for all MCP tools.
|
||||
|
||||
Manages tool registration, retrieval, and validation.
|
||||
Ensures unique tool names and provides filtering by plugin type.
|
||||
|
||||
Examples:
|
||||
>>> registry = ToolRegistry()
|
||||
>>> tool = ToolDefinition(
|
||||
... name="wordpress_list_posts",
|
||||
... description="List WordPress posts",
|
||||
... handler=async_handler_func,
|
||||
... plugin_type="wordpress"
|
||||
... )
|
||||
>>> registry.register(tool)
|
||||
>>> all_tools = registry.get_all()
|
||||
>>> wp_tools = registry.get_by_plugin_type("wordpress")
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize empty tool registry."""
|
||||
self.tools: dict[str, ToolDefinition] = {}
|
||||
self.logger = logging.getLogger("ToolRegistry")
|
||||
self.logger.info("ToolRegistry initialized")
|
||||
|
||||
def register(self, tool: ToolDefinition) -> None:
|
||||
"""
|
||||
Register a tool in the registry.
|
||||
|
||||
Args:
|
||||
tool: Tool definition to register
|
||||
|
||||
Raises:
|
||||
ValueError: If tool name already exists
|
||||
|
||||
Examples:
|
||||
>>> registry.register(tool_definition)
|
||||
"""
|
||||
if tool.name in self.tools:
|
||||
raise ValueError(f"Tool '{tool.name}' already registered")
|
||||
|
||||
self.tools[tool.name] = tool
|
||||
self.logger.debug(f"Registered tool: {tool.name} ({tool.plugin_type})")
|
||||
|
||||
def register_many(self, tools: list[ToolDefinition]) -> int:
|
||||
"""
|
||||
Register multiple tools at once.
|
||||
|
||||
Args:
|
||||
tools: List of tool definitions
|
||||
|
||||
Returns:
|
||||
Number of tools successfully registered
|
||||
|
||||
Examples:
|
||||
>>> count = registry.register_many([tool1, tool2, tool3])
|
||||
>>> print(f"Registered {count} tools")
|
||||
"""
|
||||
count = 0
|
||||
for tool in tools:
|
||||
try:
|
||||
self.register(tool)
|
||||
count += 1
|
||||
except ValueError as e:
|
||||
self.logger.warning(f"Skipped duplicate tool: {e}")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to register tool: {e}", exc_info=True)
|
||||
|
||||
self.logger.info(f"Registered {count}/{len(tools)} tools")
|
||||
return count
|
||||
|
||||
def get_all(self) -> list[ToolDefinition]:
|
||||
"""
|
||||
Get all registered tools.
|
||||
|
||||
Returns:
|
||||
List of all tool definitions
|
||||
|
||||
Examples:
|
||||
>>> tools = registry.get_all()
|
||||
>>> print(f"Total tools: {len(tools)}")
|
||||
"""
|
||||
return list(self.tools.values())
|
||||
|
||||
def get_by_name(self, name: str) -> ToolDefinition | None:
|
||||
"""
|
||||
Get a tool by its name.
|
||||
|
||||
Args:
|
||||
name: Tool name
|
||||
|
||||
Returns:
|
||||
Tool definition if found, None otherwise
|
||||
|
||||
Examples:
|
||||
>>> tool = registry.get_by_name("wordpress_get_post")
|
||||
>>> if tool:
|
||||
... print(f"Found: {tool.description}")
|
||||
"""
|
||||
return self.tools.get(name)
|
||||
|
||||
def get_by_plugin_type(self, plugin_type: str) -> list[ToolDefinition]:
|
||||
"""
|
||||
Get all tools for a specific plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type (e.g., "wordpress", "gitea")
|
||||
|
||||
Returns:
|
||||
List of tools for the specified plugin type
|
||||
|
||||
Examples:
|
||||
>>> wp_tools = registry.get_by_plugin_type("wordpress")
|
||||
>>> print(f"WordPress tools: {len(wp_tools)}")
|
||||
"""
|
||||
return [tool for tool in self.tools.values() if tool.plugin_type == plugin_type]
|
||||
|
||||
def get_count(self) -> int:
|
||||
"""
|
||||
Get total number of registered tools.
|
||||
|
||||
Returns:
|
||||
Count of registered tools
|
||||
|
||||
Examples:
|
||||
>>> count = registry.get_count()
|
||||
"""
|
||||
return len(self.tools)
|
||||
|
||||
def get_count_by_plugin(self) -> dict[str, int]:
|
||||
"""
|
||||
Get tool counts grouped by plugin type.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping plugin type to tool count
|
||||
|
||||
Examples:
|
||||
>>> counts = registry.get_count_by_plugin()
|
||||
>>> print(counts) # {'wordpress': 95, 'gitea': 50}
|
||||
"""
|
||||
counts = {}
|
||||
for tool in self.tools.values():
|
||||
plugin_type = tool.plugin_type
|
||||
counts[plugin_type] = counts.get(plugin_type, 0) + 1
|
||||
return counts
|
||||
|
||||
def clear(self) -> None:
|
||||
"""
|
||||
Clear all registered tools.
|
||||
|
||||
Primarily for testing purposes.
|
||||
|
||||
Examples:
|
||||
>>> registry.clear()
|
||||
>>> assert registry.get_count() == 0
|
||||
"""
|
||||
self.tools.clear()
|
||||
self.logger.info("Tool registry cleared")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""String representation of registry."""
|
||||
counts = self.get_count_by_plugin()
|
||||
counts_str = ", ".join(f"{k}: {v}" for k, v in counts.items())
|
||||
return f"ToolRegistry(total={self.get_count()}, {counts_str})"
|
||||
|
||||
# Singleton instance
|
||||
_tool_registry: ToolRegistry | None = None
|
||||
|
||||
def get_tool_registry() -> ToolRegistry:
|
||||
"""
|
||||
Get the singleton tool registry instance.
|
||||
|
||||
Returns:
|
||||
Global ToolRegistry instance
|
||||
|
||||
Examples:
|
||||
>>> registry = get_tool_registry()
|
||||
>>> registry.register(tool)
|
||||
"""
|
||||
global _tool_registry
|
||||
if _tool_registry is None:
|
||||
_tool_registry = ToolRegistry()
|
||||
return _tool_registry
|
||||
362
core/unified_tools.py
Normal file
362
core/unified_tools.py
Normal file
@@ -0,0 +1,362 @@
|
||||
"""
|
||||
Unified Tool Generator
|
||||
|
||||
Generates context-based tools that work across multiple sites.
|
||||
Maintains backward compatibility by keeping per-site tools.
|
||||
|
||||
Architecture:
|
||||
- Old: wordpress_site1_get_post(post_id)
|
||||
- New: wordpress_get_post(site, post_id)
|
||||
- Both work simultaneously!
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from core.site_registry import get_site_registry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class UnifiedToolGenerator:
|
||||
"""
|
||||
Generates unified tools from per-site tool definitions.
|
||||
|
||||
Takes existing plugin tools and creates context-based versions
|
||||
that accept a 'site' parameter for dynamic routing.
|
||||
"""
|
||||
|
||||
def __init__(self, project_manager):
|
||||
"""
|
||||
Initialize unified tool generator.
|
||||
|
||||
Args:
|
||||
project_manager: ProjectManager instance with discovered projects
|
||||
"""
|
||||
self.project_manager = project_manager
|
||||
self.site_registry = get_site_registry()
|
||||
self.logger = logging.getLogger("UnifiedToolGenerator")
|
||||
|
||||
def generate_unified_tools(self, plugin_type: str) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Generate unified tools for a specific plugin type.
|
||||
|
||||
Args:
|
||||
plugin_type: Type of plugin (e.g., 'wordpress')
|
||||
|
||||
Returns:
|
||||
List of unified tool definitions
|
||||
"""
|
||||
# Get all projects of this type
|
||||
projects = self.project_manager.get_projects_by_type(plugin_type)
|
||||
|
||||
if not projects:
|
||||
self.logger.warning(f"No projects found for plugin type: {plugin_type}")
|
||||
return []
|
||||
|
||||
# Use the first project as a template to get tool definitions
|
||||
first_project_id = list(projects.keys())[0]
|
||||
template_plugin = projects[first_project_id]
|
||||
template_tools = template_plugin.get_tools()
|
||||
|
||||
self.logger.info(
|
||||
f"Generating unified tools for {plugin_type} "
|
||||
f"from {len(template_tools)} template tools"
|
||||
)
|
||||
|
||||
unified_tools = []
|
||||
seen_actions = set()
|
||||
|
||||
for tool in template_tools:
|
||||
# Extract action name from per-site tool name
|
||||
# e.g., "wordpress_site1_get_post" -> "get_post"
|
||||
tool_name = tool["name"]
|
||||
parts = tool_name.split("_")
|
||||
|
||||
# Skip if not in expected format
|
||||
if len(parts) < 3:
|
||||
continue
|
||||
|
||||
# Extract action (everything after plugin_type_site_id_)
|
||||
# e.g., wordpress_site1_get_post -> get_post
|
||||
action = "_".join(parts[2:])
|
||||
|
||||
# Skip duplicates (we only need one unified tool per action)
|
||||
if action in seen_actions:
|
||||
continue
|
||||
seen_actions.add(action)
|
||||
|
||||
# Create unified tool
|
||||
unified_tool = self._create_unified_tool(
|
||||
plugin_type=plugin_type, action=action, template_tool=tool
|
||||
)
|
||||
|
||||
if unified_tool:
|
||||
unified_tools.append(unified_tool)
|
||||
|
||||
self.logger.info(f"Generated {len(unified_tools)} unified tools for {plugin_type}")
|
||||
return unified_tools
|
||||
|
||||
def _create_unified_tool(
|
||||
self, plugin_type: str, action: str, template_tool: dict[str, Any]
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Create a unified tool from a template.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type (e.g., 'wordpress')
|
||||
action: Action name (e.g., 'get_post')
|
||||
template_tool: Original per-site tool definition
|
||||
|
||||
Returns:
|
||||
Unified tool definition
|
||||
"""
|
||||
try:
|
||||
# Create unified tool name
|
||||
unified_name = f"{plugin_type}_{action}"
|
||||
|
||||
# Get available sites for this plugin type
|
||||
site_options = self.site_registry.get_site_options(plugin_type)
|
||||
|
||||
if not site_options:
|
||||
self.logger.warning(f"No sites available for {plugin_type}, skipping {action}")
|
||||
return None
|
||||
|
||||
# Modify input schema to add 'site' parameter
|
||||
original_schema = template_tool.get("inputSchema", {})
|
||||
unified_schema = self._add_site_parameter(original_schema, plugin_type, site_options)
|
||||
|
||||
# Update description to mention site parameter
|
||||
original_description = template_tool.get("description", "")
|
||||
unified_description = self._update_description(original_description, plugin_type)
|
||||
|
||||
# Create wrapper handler
|
||||
original_handler = template_tool.get("handler")
|
||||
unified_handler = self._create_unified_handler(plugin_type, action, original_handler)
|
||||
|
||||
return {
|
||||
"name": unified_name,
|
||||
"description": unified_description,
|
||||
"inputSchema": unified_schema,
|
||||
"handler": unified_handler,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Error creating unified tool for {plugin_type}_{action}: {e}", exc_info=True
|
||||
)
|
||||
return None
|
||||
|
||||
def _add_site_parameter(
|
||||
self, original_schema: dict[str, Any], plugin_type: str, site_options: list[str]
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Add 'site' parameter to input schema.
|
||||
|
||||
Args:
|
||||
original_schema: Original input schema
|
||||
plugin_type: Plugin type
|
||||
site_options: Available site IDs/aliases
|
||||
|
||||
Returns:
|
||||
Modified schema with site parameter
|
||||
"""
|
||||
# Deep copy to avoid modifying original
|
||||
import copy
|
||||
|
||||
schema = copy.deepcopy(original_schema)
|
||||
|
||||
# Ensure schema has required structure
|
||||
if "properties" not in schema:
|
||||
schema["properties"] = {}
|
||||
if "required" not in schema:
|
||||
schema["required"] = []
|
||||
|
||||
# Add 'site' as first parameter
|
||||
schema["properties"] = {
|
||||
"site": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
f"Site ID or alias. Available options: {', '.join(site_options)}. "
|
||||
f"Use list_projects() to see all configured sites."
|
||||
),
|
||||
"enum": site_options,
|
||||
},
|
||||
**schema["properties"],
|
||||
}
|
||||
|
||||
# Make 'site' required
|
||||
if "site" not in schema["required"]:
|
||||
schema["required"].insert(0, "site")
|
||||
|
||||
return schema
|
||||
|
||||
def _update_description(self, original_description: str, plugin_type: str) -> str:
|
||||
"""
|
||||
Update tool description to mention unified context.
|
||||
|
||||
Args:
|
||||
original_description: Original description
|
||||
plugin_type: Plugin type
|
||||
|
||||
Returns:
|
||||
Updated description
|
||||
"""
|
||||
# Remove site-specific mentions (e.g., "from site1", "in site2")
|
||||
import re
|
||||
|
||||
cleaned = re.sub(
|
||||
r"\b(?:from|in|for)\s+site\d+\b", "", original_description, flags=re.IGNORECASE
|
||||
)
|
||||
cleaned = re.sub(
|
||||
r"\bsite\d+\b", f"the specified {plugin_type} site", cleaned, flags=re.IGNORECASE
|
||||
)
|
||||
|
||||
# Add unified context note
|
||||
prefix = "[UNIFIED] "
|
||||
if not cleaned.startswith(prefix):
|
||||
cleaned = prefix + cleaned
|
||||
|
||||
return cleaned.strip()
|
||||
|
||||
def _create_unified_handler(
|
||||
self, plugin_type: str, action: str, original_handler: Callable | None
|
||||
) -> Callable:
|
||||
"""
|
||||
Create a unified handler that routes to the correct site.
|
||||
|
||||
Args:
|
||||
plugin_type: Plugin type
|
||||
action: Action name
|
||||
original_handler: Original handler (not used, we call plugin method directly)
|
||||
|
||||
Returns:
|
||||
Async handler function
|
||||
"""
|
||||
|
||||
async def unified_handler(site: str, **kwargs):
|
||||
"""
|
||||
Unified handler that routes to the correct site plugin.
|
||||
|
||||
Args:
|
||||
site: Site ID or alias
|
||||
**kwargs: Other parameters for the tool
|
||||
"""
|
||||
try:
|
||||
# Get site info from registry
|
||||
site_info = self.site_registry.get_site(plugin_type, site)
|
||||
|
||||
if not site_info:
|
||||
available = self.site_registry.get_site_options(plugin_type)
|
||||
return (
|
||||
f"Error: Site '{site}' not found for {plugin_type}. "
|
||||
f"Available sites: {', '.join(available)}"
|
||||
)
|
||||
|
||||
# Get the plugin instance
|
||||
full_id = site_info.get_full_id()
|
||||
|
||||
# SECURITY: Check if API key has access to this project
|
||||
from core.context import get_api_key_context
|
||||
|
||||
api_key_info = get_api_key_context()
|
||||
|
||||
if api_key_info and not api_key_info.get("is_global"):
|
||||
# Per-project key - must match the project
|
||||
allowed_project = api_key_info.get("project_id")
|
||||
|
||||
# Resolve allowed_project to normalize alias vs site_id
|
||||
# API key might have been created with alias (wordpress_myblog)
|
||||
# or site_id (wordpress_site1)
|
||||
allowed_project_normalized = allowed_project
|
||||
if allowed_project and "_" in allowed_project:
|
||||
# Extract plugin type and site identifier from allowed_project
|
||||
allowed_parts = allowed_project.split("_", 1)
|
||||
if len(allowed_parts) == 2:
|
||||
allowed_plugin_type, allowed_site_identifier = allowed_parts
|
||||
# Try to resolve the site identifier to site_id
|
||||
try:
|
||||
allowed_site_info = self.site_registry.get_site(
|
||||
allowed_plugin_type, allowed_site_identifier
|
||||
)
|
||||
if allowed_site_info:
|
||||
# Normalize to plugin_type_site_id format
|
||||
allowed_project_normalized = allowed_site_info.get_full_id()
|
||||
except (ValueError, Exception):
|
||||
# Site not found, keep original for error message
|
||||
pass
|
||||
|
||||
if allowed_project_normalized != full_id:
|
||||
logger.warning(
|
||||
f"Access denied: API key for project '{allowed_project}' "
|
||||
f"attempted to access '{full_id}'"
|
||||
)
|
||||
return (
|
||||
f"Error: Access denied. This API key is restricted to project '{allowed_project}'. "
|
||||
f"Use a global API key or create a key for '{full_id}'."
|
||||
)
|
||||
|
||||
plugin = self.project_manager.get_project(full_id)
|
||||
|
||||
if not plugin:
|
||||
return f"Error: Plugin instance not found for {full_id}"
|
||||
|
||||
# Find the original handler method in the plugin
|
||||
# The original per-site tool name was: {plugin_type}_{site_id}_{action}
|
||||
original_tool_name = f"{plugin_type}_{site_info.site_id}_{action}"
|
||||
|
||||
# Get all tools from plugin and find the matching handler
|
||||
tools = plugin.get_tools()
|
||||
handler = None
|
||||
|
||||
for tool in tools:
|
||||
if tool["name"] == original_tool_name:
|
||||
handler = tool.get("handler")
|
||||
break
|
||||
|
||||
if not handler:
|
||||
return (
|
||||
f"Error: Handler not found for {original_tool_name}. "
|
||||
f"This may be a plugin implementation issue."
|
||||
)
|
||||
|
||||
# Filter out None values from kwargs to avoid validation errors
|
||||
# WordPress API doesn't accept None values in query parameters
|
||||
filtered_kwargs = {key: value for key, value in kwargs.items() if value is not None}
|
||||
|
||||
# Call the handler with filtered kwargs
|
||||
result = await handler(**filtered_kwargs)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error in unified handler for {plugin_type}_{action}: {e}", exc_info=True
|
||||
)
|
||||
return f"Error: {str(e)}"
|
||||
|
||||
return unified_handler
|
||||
|
||||
def generate_all_unified_tools(self) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Generate unified tools for all registered plugin types.
|
||||
|
||||
Returns:
|
||||
List of all unified tool definitions
|
||||
"""
|
||||
all_tools = []
|
||||
|
||||
# Get all plugin types from registry
|
||||
from plugins import registry
|
||||
|
||||
plugin_types = registry.get_registered_types()
|
||||
|
||||
for plugin_type in plugin_types:
|
||||
tools = self.generate_unified_tools(plugin_type)
|
||||
all_tools.extend(tools)
|
||||
|
||||
self.logger.info(
|
||||
f"Generated {len(all_tools)} total unified tools "
|
||||
f"across {len(plugin_types)} plugin types"
|
||||
)
|
||||
|
||||
return all_tools
|
||||
Reference in New Issue
Block a user