""" API Server for PyGuardian Controller REST API endpoints for agent authentication and cluster management """ import json import logging import asyncio from datetime import datetime from typing import Dict, Any, Optional, Tuple from aiohttp import web, WSMsgType from aiohttp.web import Application, Request, Response, WebSocketResponse import aiohttp_cors import ssl from pathlib import Path from .auth import AgentAuthentication, AgentAuthenticationError from .cluster_manager import ClusterManager from .storage import Storage logger = logging.getLogger(__name__) class PyGuardianAPI: """ PyGuardian Controller API Server Provides REST API and WebSocket endpoints for agent communication """ def __init__(self, cluster_manager: ClusterManager, config: Dict[str, Any]): self.cluster_manager = cluster_manager self.config = config self.app = None self.server = None self.websockets = set() # Active WebSocket connections # API configuration self.host = config.get('api_host', '0.0.0.0') self.port = config.get('api_port', 8443) self.ssl_cert = config.get('ssl_cert') self.ssl_key = config.get('ssl_key') self.api_secret = config.get('api_secret', 'change-this-secret') async def create_app(self) -> Application: """Create aiohttp application with routes and middleware""" app = web.Application() # Add CORS support cors = aiohttp_cors.setup(app, defaults={ "*": aiohttp_cors.ResourceOptions( allow_credentials=True, expose_headers="*", allow_headers="*", allow_methods="*" ) }) # Add routes self._setup_routes(app) # Add CORS to all routes for route in list(app.router.routes()): cors.add(route) # Add middleware app.middlewares.append(self._auth_middleware) app.middlewares.append(self._error_middleware) self.app = app return app def _setup_routes(self, app: Application): """Setup API routes""" # Health check app.router.add_get('/health', self.health_check) # Agent authentication endpoints app.router.add_post('/api/v1/auth/register', self.register_agent) app.router.add_post('/api/v1/auth/login', self.login_agent) app.router.add_post('/api/v1/auth/refresh', self.refresh_token) app.router.add_post('/api/v1/auth/logout', self.logout_agent) app.router.add_post('/api/v1/auth/verify', self.verify_token) # Cluster management endpoints app.router.add_get('/api/v1/cluster/status', self.cluster_status) app.router.add_get('/api/v1/cluster/agents', self.list_agents) app.router.add_get('/api/v1/cluster/agents/{agent_id}', self.get_agent_info) app.router.add_post('/api/v1/cluster/agents/{agent_id}/deploy', self.deploy_agent) app.router.add_delete('/api/v1/cluster/agents/{agent_id}', self.remove_agent) # Agent communication endpoints app.router.add_post('/api/v1/agent/heartbeat', self.agent_heartbeat) app.router.add_post('/api/v1/agent/report', self.agent_report) app.router.add_get('/api/v1/agent/config', self.get_agent_config) app.router.add_post('/api/v1/agent/logs', self.upload_agent_logs) # WebSocket endpoint for real-time communication app.router.add_get('/ws/agent', self.websocket_handler) # Metrics endpoint for monitoring app.router.add_get('/metrics', self.metrics_endpoint) async def _auth_middleware(self, request: Request, handler): """Authentication middleware for protected endpoints""" # Skip auth for health check and public endpoints if request.path in ['/health', '/metrics'] or request.path.startswith('/api/v1/auth/'): return await handler(request) # Check for API secret (for controller-to-controller communication) api_secret = request.headers.get('X-API-Secret') if api_secret and api_secret == self.api_secret: request['authenticated'] = True request['auth_type'] = 'api_secret' return await handler(request) # Check for agent token auth_header = request.headers.get('Authorization', '') if not auth_header.startswith('Bearer '): return web.json_response( {'error': 'Missing or invalid authorization header'}, status=401 ) token = auth_header[7:] # Remove 'Bearer ' prefix try: success, agent_id = await self.cluster_manager.verify_agent_token(token) if success: request['authenticated'] = True request['auth_type'] = 'agent_token' request['agent_id'] = agent_id return await handler(request) else: return web.json_response( {'error': 'Invalid or expired token'}, status=401 ) except Exception as e: logger.error(f"Authentication error: {e}") return web.json_response( {'error': 'Authentication failed'}, status=401 ) async def _error_middleware(self, request: Request, handler): """Error handling middleware""" try: return await handler(request) except web.HTTPException: raise except Exception as e: logger.error(f"API error in {request.path}: {e}") return web.json_response( {'error': 'Internal server error'}, status=500 ) # ========================= # API Endpoint Handlers # ========================= async def health_check(self, request: Request) -> Response: """Health check endpoint""" return web.json_response({ 'status': 'healthy', 'timestamp': datetime.now().isoformat(), 'version': '2.0.0', 'cluster': self.cluster_manager.cluster_name }) async def register_agent(self, request: Request) -> Response: """Register new agent endpoint""" try: data = await request.json() # Validate required fields required_fields = ['hostname', 'ip_address'] for field in required_fields: if field not in data: return web.json_response( {'error': f'Missing required field: {field}'}, status=400 ) # Register agent success, result = await self.cluster_manager.register_new_agent( hostname=data['hostname'], ip_address=data['ip_address'], ssh_user=data.get('ssh_user', 'root'), ssh_port=data.get('ssh_port', 22), ssh_key_path=data.get('ssh_key_path'), ssh_password=data.get('ssh_password') ) if success: # Don't return sensitive data in logs safe_result = {k: v for k, v in result.items() if k not in ['secret_key']} logger.info(f"Registered new agent: {safe_result}") return web.json_response(result, status=201) else: return web.json_response(result, status=400) except Exception as e: logger.error(f"Agent registration error: {e}") return web.json_response( {'error': 'Registration failed'}, status=500 ) async def login_agent(self, request: Request) -> Response: """Agent login endpoint""" try: data = await request.json() # Validate required fields if 'agent_id' not in data or 'secret_key' not in data: return web.json_response( {'error': 'Missing agent_id or secret_key'}, status=400 ) client_ip = request.remote or 'unknown' # Authenticate agent success, result = await self.cluster_manager.authenticate_agent( agent_id=data['agent_id'], secret_key=data['secret_key'], ip_address=client_ip ) if success: logger.info(f"Agent {data['agent_id']} authenticated from {client_ip}") return web.json_response(result) else: return web.json_response(result, status=401) except Exception as e: logger.error(f"Agent login error: {e}") return web.json_response( {'error': 'Authentication failed'}, status=500 ) async def refresh_token(self, request: Request) -> Response: """Token refresh endpoint""" try: data = await request.json() if 'refresh_token' not in data: return web.json_response( {'error': 'Missing refresh_token'}, status=400 ) success, result = await self.cluster_manager.refresh_agent_token( data['refresh_token'] ) if success: return web.json_response(result) else: return web.json_response(result, status=401) except Exception as e: logger.error(f"Token refresh error: {e}") return web.json_response( {'error': 'Token refresh failed'}, status=500 ) async def logout_agent(self, request: Request) -> Response: """Agent logout endpoint""" agent_id = request.get('agent_id') if not agent_id: return web.json_response( {'error': 'Agent ID not found'}, status=400 ) try: success = await self.cluster_manager.revoke_agent_access(agent_id) if success: return web.json_response({'message': 'Logged out successfully'}) else: return web.json_response( {'error': 'Logout failed'}, status=500 ) except Exception as e: logger.error(f"Agent logout error: {e}") return web.json_response( {'error': 'Logout failed'}, status=500 ) async def verify_token(self, request: Request) -> Response: """Token verification endpoint""" try: data = await request.json() if 'token' not in data: return web.json_response( {'error': 'Missing token'}, status=400 ) success, agent_id = await self.cluster_manager.verify_agent_token( data['token'] ) if success: return web.json_response({ 'valid': True, 'agent_id': agent_id }) else: return web.json_response({ 'valid': False, 'error': agent_id # agent_id contains error message on failure }) except Exception as e: logger.error(f"Token verification error: {e}") return web.json_response( {'error': 'Verification failed'}, status=500 ) async def cluster_status(self, request: Request) -> Response: """Get cluster status""" try: status = await self.cluster_manager.get_cluster_stats() auth_status = await self.cluster_manager.get_cluster_auth_status() return web.json_response({ 'cluster_info': status, 'authentication': auth_status }) except Exception as e: logger.error(f"Cluster status error: {e}") return web.json_response( {'error': 'Failed to get cluster status'}, status=500 ) async def list_agents(self, request: Request) -> Response: """List all agents in cluster""" try: agents = await self.cluster_manager.get_cluster_agents() return web.json_response({'agents': agents}) except Exception as e: logger.error(f"List agents error: {e}") return web.json_response( {'error': 'Failed to list agents'}, status=500 ) async def get_agent_info(self, request: Request) -> Response: """Get specific agent information""" agent_id = request.match_info['agent_id'] try: if agent_id in self.cluster_manager.agents: agent_info = self.cluster_manager.agents[agent_id].to_dict() # Get additional info auth_logs = await self.cluster_manager.get_agent_auth_logs(agent_id) sessions = await self.cluster_manager.get_active_agent_sessions(agent_id) return web.json_response({ 'agent': agent_info, 'auth_logs': auth_logs[:10], # Last 10 logs 'sessions': sessions }) else: return web.json_response( {'error': 'Agent not found'}, status=404 ) except Exception as e: logger.error(f"Get agent info error: {e}") return web.json_response( {'error': 'Failed to get agent info'}, status=500 ) async def deploy_agent(self, request: Request) -> Response: """Deploy agent endpoint""" agent_id = request.match_info['agent_id'] try: data = await request.json() force_reinstall = data.get('force_reinstall', False) success, message = await self.cluster_manager.deploy_agent( agent_id, force_reinstall ) if success: return web.json_response({ 'success': True, 'message': message }) else: return web.json_response({ 'success': False, 'message': message }, status=400) except Exception as e: logger.error(f"Deploy agent error: {e}") return web.json_response( {'error': 'Deployment failed'}, status=500 ) async def remove_agent(self, request: Request) -> Response: """Remove agent endpoint""" agent_id = request.match_info['agent_id'] try: # Parse query parameters cleanup_remote = request.query.get('cleanup_remote', 'false').lower() == 'true' success, message = await self.cluster_manager.remove_agent( agent_id, cleanup_remote ) if success: return web.json_response({ 'success': True, 'message': message }) else: return web.json_response({ 'success': False, 'message': message }, status=400) except Exception as e: logger.error(f"Remove agent error: {e}") return web.json_response( {'error': 'Agent removal failed'}, status=500 ) async def agent_heartbeat(self, request: Request) -> Response: """Agent heartbeat endpoint""" agent_id = request.get('agent_id') if not agent_id: return web.json_response( {'error': 'Agent ID not found'}, status=400 ) try: data = await request.json() # Update agent status if agent_id in self.cluster_manager.agents: agent = self.cluster_manager.agents[agent_id] agent.last_check = datetime.now() agent.status = 'online' agent.stats.update(data.get('stats', {})) # Send any pending commands return web.json_response({ 'status': 'ok', 'next_heartbeat': 60, # seconds 'commands': [] # TODO: implement command queue }) except Exception as e: logger.error(f"Heartbeat error: {e}") return web.json_response( {'error': 'Heartbeat failed'}, status=500 ) async def agent_report(self, request: Request) -> Response: """Agent security report endpoint""" agent_id = request.get('agent_id') if not agent_id: return web.json_response( {'error': 'Agent ID not found'}, status=400 ) try: report = await request.json() # Process security report logger.info(f"Received security report from agent {agent_id}") # TODO: Process and store security events return web.json_response({'status': 'received'}) except Exception as e: logger.error(f"Agent report error: {e}") return web.json_response( {'error': 'Report processing failed'}, status=500 ) async def get_agent_config(self, request: Request) -> Response: """Get agent configuration""" agent_id = request.get('agent_id') if not agent_id: return web.json_response( {'error': 'Agent ID not found'}, status=400 ) try: # Return agent-specific configuration config = { 'heartbeat_interval': 60, 'report_interval': 300, 'log_level': 'INFO', 'features': { 'firewall_monitoring': True, 'intrusion_detection': True, 'log_analysis': True } } return web.json_response(config) except Exception as e: logger.error(f"Get agent config error: {e}") return web.json_response( {'error': 'Config retrieval failed'}, status=500 ) async def upload_agent_logs(self, request: Request) -> Response: """Upload agent logs endpoint""" agent_id = request.get('agent_id') if not agent_id: return web.json_response( {'error': 'Agent ID not found'}, status=400 ) try: logs = await request.json() # Process and store logs logger.info(f"Received {len(logs.get('entries', []))} log entries from agent {agent_id}") # TODO: Store logs in database or forward to log aggregator return web.json_response({'status': 'received'}) except Exception as e: logger.error(f"Log upload error: {e}") return web.json_response( {'error': 'Log upload failed'}, status=500 ) async def websocket_handler(self, request: Request) -> WebSocketResponse: """WebSocket endpoint for real-time agent communication""" ws = web.WebSocketResponse() await ws.prepare(request) agent_id = None try: # Add to active connections self.websockets.add(ws) async for msg in ws: if msg.type == WSMsgType.TEXT: try: data = json.loads(msg.data) if data.get('type') == 'auth' and not agent_id: # Authenticate WebSocket connection token = data.get('token') if token: success, agent_id = await self.cluster_manager.verify_agent_token(token) if success: await ws.send_text(json.dumps({ 'type': 'auth_success', 'agent_id': agent_id })) else: await ws.send_text(json.dumps({ 'type': 'auth_failed', 'error': 'Invalid token' })) elif agent_id: # Handle authenticated messages await self._handle_ws_message(ws, agent_id, data) else: await ws.send_text(json.dumps({ 'type': 'error', 'error': 'Not authenticated' })) except json.JSONDecodeError: await ws.send_text(json.dumps({ 'type': 'error', 'error': 'Invalid JSON' })) elif msg.type == WSMsgType.ERROR: logger.error(f'WebSocket error: {ws.exception()}') except Exception as e: logger.error(f"WebSocket error: {e}") finally: self.websockets.discard(ws) return ws async def _handle_ws_message(self, ws: WebSocketResponse, agent_id: str, data: Dict[str, Any]): """Handle authenticated WebSocket message""" message_type = data.get('type') if message_type == 'ping': await ws.send_text(json.dumps({'type': 'pong'})) elif message_type == 'status_update': # Handle agent status update if agent_id in self.cluster_manager.agents: agent = self.cluster_manager.agents[agent_id] agent.status = data.get('status', 'unknown') agent.last_check = datetime.now() else: await ws.send_text(json.dumps({ 'type': 'error', 'error': f'Unknown message type: {message_type}' })) async def metrics_endpoint(self, request: Request) -> Response: """Prometheus metrics endpoint""" try: metrics = [] # Cluster metrics stats = await self.cluster_manager.get_cluster_stats() auth_status = await self.cluster_manager.get_cluster_auth_status() metrics.append(f"pyguardian_cluster_total_agents {stats['total_agents']}") metrics.append(f"pyguardian_cluster_online_agents {stats['online_agents']}") metrics.append(f"pyguardian_cluster_offline_agents {stats['offline_agents']}") metrics.append(f"pyguardian_cluster_deployed_agents {stats['deployed_agents']}") metrics.append(f"pyguardian_cluster_authenticated_agents {auth_status['authenticated_agents']}") metrics.append(f"pyguardian_cluster_unauthenticated_agents {auth_status['unauthenticated_agents']}") # WebSocket connections metrics.append(f"pyguardian_websocket_connections {len(self.websockets)}") return web.Response( text='\\n'.join(metrics), content_type='text/plain' ) except Exception as e: logger.error(f"Metrics error: {e}") return web.Response( text='# Error generating metrics', content_type='text/plain', status=500 ) # ========================= # Server Management # ========================= async def start_server(self): """Start the API server""" try: app = await self.create_app() # Setup SSL context if certificates are provided ssl_context = None if self.ssl_cert and self.ssl_key: ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) ssl_context.load_cert_chain(self.ssl_cert, self.ssl_key) logger.info("SSL enabled for API server") # Start server runner = web.AppRunner(app) await runner.setup() site = web.TCPSite(runner, self.host, self.port, ssl_context=ssl_context) await site.start() protocol = 'https' if ssl_context else 'http' logger.info(f"PyGuardian API server started on {protocol}://{self.host}:{self.port}") self.server = runner except Exception as e: logger.error(f"Failed to start API server: {e}") raise async def stop_server(self): """Stop the API server""" if self.server: await self.server.cleanup() self.server = None logger.info("API server stopped") async def broadcast_to_agents(self, message: Dict[str, Any]): """Broadcast message to all connected agents via WebSocket""" if not self.websockets: return message_text = json.dumps(message) disconnected = set() for ws in self.websockets: try: await ws.send_text(message_text) except Exception: disconnected.add(ws) # Clean up disconnected WebSockets self.websockets -= disconnected