315 lines
9.9 KiB
Python
315 lines
9.9 KiB
Python
"""Environment builder orchestration with security validation."""
|
|
|
|
import re
|
|
import shutil
|
|
import threading
|
|
from typing import List, Optional, Dict, Set
|
|
from pathlib import Path
|
|
from datetime import datetime, timedelta
|
|
|
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
|
from mcp_forge.podman.client import PodmanClient
|
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
|
from mcp_forge.builder.package_validator import PackageValidator
|
|
from mcp_forge.builder.uv_installer import UVInstaller
|
|
from mcp_forge.builder.image_builder import ImageBuilder, BuildResult
|
|
|
|
|
|
class BuildRateLimiter:
|
|
"""Rate limiter for build requests."""
|
|
|
|
def __init__(self, max_requests: int, period_seconds: int):
|
|
"""
|
|
Initialize rate limiter.
|
|
|
|
Args:
|
|
max_requests: Maximum requests per period
|
|
period_seconds: Period length in seconds
|
|
"""
|
|
self.max_requests = max_requests
|
|
self.period_seconds = period_seconds
|
|
self.requests: Dict[str, List[datetime]] = {}
|
|
self.lock = threading.Lock()
|
|
|
|
def check_rate_limit(self, user_id: str) -> None:
|
|
"""
|
|
Check if user is within rate limit.
|
|
|
|
Args:
|
|
user_id: User identifier
|
|
|
|
Raises:
|
|
RuntimeError: If rate limit exceeded
|
|
"""
|
|
with self.lock:
|
|
now = datetime.now()
|
|
|
|
# Initialize user's request list if needed
|
|
if user_id not in self.requests:
|
|
self.requests[user_id] = []
|
|
|
|
# Clean up old requests
|
|
self._cleanup_old_requests(user_id, now)
|
|
|
|
# Check if at limit
|
|
if len(self.requests[user_id]) >= self.max_requests:
|
|
raise RuntimeError(
|
|
f"Rate limit exceeded: {self.max_requests} requests "
|
|
f"per {self.period_seconds} seconds"
|
|
)
|
|
|
|
# Record this request
|
|
self.requests[user_id].append(now)
|
|
|
|
def _cleanup_old_requests(self, user_id: str, now: datetime) -> None:
|
|
"""
|
|
Remove requests older than period.
|
|
|
|
Args:
|
|
user_id: User identifier
|
|
now: Current time
|
|
"""
|
|
cutoff = now - timedelta(seconds=self.period_seconds)
|
|
self.requests[user_id] = [
|
|
req_time for req_time in self.requests[user_id]
|
|
if req_time > cutoff
|
|
]
|
|
|
|
|
|
class EnvironmentBuilder:
|
|
"""Builds custom Python environments with security validation."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: EnvironmentBuilderConfig,
|
|
podman_client: PodmanClient,
|
|
audit_logger: AuditLogger
|
|
):
|
|
"""
|
|
Initialize environment builder.
|
|
|
|
Args:
|
|
config: Environment builder configuration
|
|
podman_client: Podman API client
|
|
audit_logger: Audit logging instance
|
|
"""
|
|
self.config = config
|
|
self.podman = podman_client
|
|
self.audit_logger = audit_logger
|
|
|
|
# Initialize sub-components
|
|
self.package_validator = PackageValidator(
|
|
config.package_validation
|
|
)
|
|
self.uv_installer = UVInstaller(config.uv_cache_path)
|
|
self.image_builder = ImageBuilder(
|
|
podman_client, config, audit_logger
|
|
)
|
|
|
|
# Rate limiting and concurrency control
|
|
self.rate_limiter = BuildRateLimiter(
|
|
max_requests=config.build_rate_limit['requests'],
|
|
period_seconds=config.build_rate_limit['period']
|
|
)
|
|
self.active_builds: Set[str] = set()
|
|
self.active_builds_lock = threading.Lock()
|
|
|
|
def build_custom_environment(
|
|
self,
|
|
name: str,
|
|
packages: List[str],
|
|
base_image: str = "python:3.11-slim",
|
|
python_version: str = "3.11",
|
|
description: str = "",
|
|
user_id: str = "default"
|
|
) -> BuildResult:
|
|
"""
|
|
Build custom environment with packages.
|
|
|
|
Process:
|
|
1. Check rate limit
|
|
2. Check concurrent builds limit
|
|
3. Validate environment name
|
|
4. Validate package count
|
|
5. Validate package names (allowlist/blocklist)
|
|
6. Generate build context with UV
|
|
7. Build image
|
|
8. Cleanup build context
|
|
|
|
Args:
|
|
name: Environment name (alphanumeric + hyphens)
|
|
packages: List of package specifications
|
|
base_image: Base image to build from
|
|
python_version: Python version
|
|
description: Optional description
|
|
user_id: User ID for rate limiting
|
|
|
|
Returns:
|
|
BuildResult
|
|
|
|
Raises:
|
|
ValueError: If validation fails
|
|
SecurityError: If security check fails
|
|
RuntimeError: If rate limit or concurrency exceeded
|
|
"""
|
|
# Check rate limit
|
|
self.rate_limiter.check_rate_limit(user_id)
|
|
|
|
# Check concurrent builds
|
|
self._check_concurrent_builds()
|
|
|
|
# Validate environment name
|
|
if not re.match(r'^[a-zA-Z0-9-]+$', name):
|
|
raise ValueError(
|
|
f"Environment name '{name}' must contain only alphanumeric characters and hyphens"
|
|
)
|
|
|
|
# Validate package count
|
|
if len(packages) > self.config.max_packages:
|
|
raise ValueError(
|
|
f"Package count {len(packages)} exceeds maximum {self.config.max_packages}"
|
|
)
|
|
|
|
# Log build start
|
|
self.audit_logger.log(
|
|
event_type=AuditEventType.BUILD_REQUEST,
|
|
severity=AuditSeverity.INFO,
|
|
message=f"Building environment: {name}",
|
|
details={
|
|
"packages": packages,
|
|
"base_image": base_image,
|
|
"user_id": user_id
|
|
}
|
|
)
|
|
|
|
build_context = None
|
|
try:
|
|
# Register build as active
|
|
self._register_build_start(name)
|
|
|
|
# Validate packages
|
|
self.package_validator.validate_packages(packages)
|
|
|
|
# Generate build context
|
|
build_context = self.uv_installer.create_build_context(
|
|
packages=packages,
|
|
base_image=base_image,
|
|
python_version=python_version
|
|
)
|
|
|
|
# Build image
|
|
result = self.image_builder.build_image(
|
|
name=name,
|
|
build_context=build_context,
|
|
base_image=base_image,
|
|
packages=packages
|
|
)
|
|
|
|
# Log completion
|
|
self.audit_logger.log(
|
|
event_type=AuditEventType.BUILD_COMPLETE,
|
|
severity=AuditSeverity.INFO,
|
|
message=f"Environment built: {name}",
|
|
details={
|
|
"image_id": result.image_id,
|
|
"build_time": result.build_time,
|
|
"user_id": user_id
|
|
}
|
|
)
|
|
|
|
return result
|
|
|
|
finally:
|
|
# Always cleanup
|
|
self._register_build_complete(name)
|
|
if build_context and build_context.exists():
|
|
shutil.rmtree(build_context, ignore_errors=True)
|
|
|
|
def build_from_template(
|
|
self,
|
|
template_name: str,
|
|
additional_packages: Optional[List[str]] = None,
|
|
name: Optional[str] = None,
|
|
user_id: str = "default"
|
|
) -> BuildResult:
|
|
"""
|
|
Build environment from template.
|
|
|
|
Expands template packages and adds additional packages.
|
|
|
|
Args:
|
|
template_name: Name of template to use
|
|
additional_packages: Optional additional packages
|
|
name: Optional custom name (uses template name if not provided)
|
|
user_id: User ID for rate limiting
|
|
|
|
Returns:
|
|
BuildResult
|
|
|
|
Raises:
|
|
ValueError: If template not found
|
|
"""
|
|
# Validate template exists
|
|
if template_name not in self.config.templates:
|
|
raise ValueError(f"Template '{template_name}' not found")
|
|
|
|
template = self.config.templates[template_name]
|
|
|
|
# Combine template and additional packages
|
|
packages = template["packages"].copy()
|
|
if additional_packages:
|
|
packages.extend(additional_packages)
|
|
|
|
# Use template name if no custom name provided
|
|
if name is None:
|
|
name = template_name
|
|
|
|
# Build with combined package list
|
|
return self.build_custom_environment(
|
|
name=name,
|
|
packages=packages,
|
|
description=template.get("description", ""),
|
|
user_id=user_id
|
|
)
|
|
|
|
def list_templates(self) -> Dict[str, dict]:
|
|
"""
|
|
List available templates.
|
|
|
|
Returns:
|
|
Dictionary of template name to template metadata
|
|
"""
|
|
return self.config.templates
|
|
|
|
def _check_concurrent_builds(self) -> None:
|
|
"""
|
|
Check concurrent builds limit.
|
|
|
|
Raises:
|
|
RuntimeError: If at max concurrent builds
|
|
"""
|
|
with self.active_builds_lock:
|
|
if len(self.active_builds) >= self.config.max_concurrent_builds:
|
|
raise RuntimeError(
|
|
f"Maximum concurrent builds ({self.config.max_concurrent_builds}) reached"
|
|
)
|
|
|
|
def _register_build_start(self, name: str) -> None:
|
|
"""
|
|
Register build as started.
|
|
|
|
Args:
|
|
name: Environment name
|
|
"""
|
|
with self.active_builds_lock:
|
|
self.active_builds.add(name)
|
|
|
|
def _register_build_complete(self, name: str) -> None:
|
|
"""
|
|
Register build as completed.
|
|
|
|
Args:
|
|
name: Environment name
|
|
"""
|
|
with self.active_builds_lock:
|
|
self.active_builds.discard(name)
|