ai-agent/orchestrator/core/dispatcher.py

107 lines
No EOL
4.1 KiB
Python

import logging
import ssl
# TLS 1.3 Enforcement (Security Baseline #4)
SSL_CONTEXT = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
SSL_CONTEXT.minimum_version = ssl.TLSVersion.TLSv1_3
# TLS 1.3 ciphers are automatically negotiated - explicit configuration not required
import os
from typing import Dict, Optional
from collections import deque
from dataclasses import dataclass, field
from security.rbac_engine import RBACEngine, Role
@dataclass
class Task:
id: str
payload: dict
requester: str
priority: int = 1
metadata: dict = field(default_factory=dict) # Initialize empty dict for metadata
class TaskQueue:
def __init__(self):
self.queue = deque()
self.logger = logging.getLogger('TaskQueue')
# Enforce environment variable for encryption key
encryption_key = os.getenv('DISPATCHER_ENCRYPTION_KEY')
if not encryption_key:
raise RuntimeError("DISPATCHER_ENCRYPTION_KEY environment variable required")
self.rbac = RBACEngine(encryption_key.encode('utf-8'))
def add_task(self, task: Task) -> None:
"""Add task to queue with priority sorting"""
if not self._validate_permissions(task.requester, task):
raise PermissionError(
f"User {task.requester} lacks {task.metadata.get('action', 'execute')} "
f"permission for {task.metadata.get('resource', 'unknown')} resources"
)
# Insert sorted by priority (highest first)
if not self.queue:
self.queue.append(task)
else:
inserted = False
for i, t in enumerate(self.queue):
if task.priority > t.priority:
self.queue.insert(i, task)
inserted = True
break
if not inserted:
self.queue.append(task)
self.logger.info(f"Added task {task.id} from {task.requester}")
def get_next_task(self) -> Optional[Task]:
"""Retrieve next highest priority task"""
return self.queue.popleft() if self.queue else None
def _validate_permissions(self, user: str, task: Task) -> bool:
"""Validate user permissions using RBAC engine"""
# Extract resource/action from task metadata with defaults
resource = task.metadata.get('resource', 'tasks')
action = task.metadata.get('action', 'execute')
return self.rbac.validate_permission(user, resource, action)
class Dispatcher:
def __init__(self):
self.queue = TaskQueue()
self.active_tasks: Dict[str, Task] = {}
def dispatch(self) -> None:
"""Main dispatch loop"""
while self.queue.queue:
task = self.queue.get_next_task()
if task:
self._process_task(task)
def _process_task(self, task: Task) -> None:
"""Execute task and handle results"""
self.logger = logging.getLogger('TaskDispatcher')
self.active_tasks[task.id] = task
try:
if not self.queue.rbac.validate_permission(task.requester, 'tasks', 'execute'):
self.logger.error(f"Unauthorized task execution attempt by {task.requester}")
raise PermissionError("Insufficient privileges for task execution")
self.logger.info(f"Processing task {task.id}")
# Check for error simulation before decryption attempt
if task.payload.get("simulate_error"):
raise ValueError("Simulated processing error - test validation")
# Decrypt payload before processing
decrypted_payload = self.queue.rbac.decrypt_payload(task.payload)
task.payload = decrypted_payload
# Actual task processing logic would go here
return True # Success
except Exception as e:
self.logger.error(f"Task {task.id} failed: {str(e)}", exc_info=True)
return False # Failure
finally:
self.logger.info(f"Completed processing task {task.id}")
self.active_tasks.pop(task.id, None)
if __name__ == "__main__":
dispatcher = Dispatcher()
dispatcher.dispatch()