""" Agent Gateway WebSocket gateway for ServerKit agents. Handles agent connections, authentication, and message routing. """ import json import logging import time from flask import request from flask_socketio import Namespace, emit, disconnect from app import sockets_rooms as rooms from app.services.agent_registry import agent_registry from app.models.server import Server logger = logging.getLogger(__name__) # The per-IP auth rate limiter is transport-neutral state and lives with the # shared handshake door in agent_registry. Re-exported here for # callers/tests that historically reached it via the gateway module. from app.services.agent_registry import ( # noqa: F401 _AUTH_RATE_LIMIT, _AUTH_RATE_WINDOW, _auth_attempts, _check_auth_rate_limit, ) class AgentNamespace(Namespace): """ SocketIO namespace for agent connections. Agents connect to: /agent """ def on_connect(self): """Handle agent connection attempt""" # Connection is not authenticated yet # Agent must send auth message first logger.info("New connection from %s", request.remote_addr) def on_disconnect(self): """Handle agent disconnection""" sid = request.sid logger.info("Disconnect: %s", sid) agent_registry.unregister_agent(sid, reason='disconnect') def on_auth(self, data): """ Handle agent authentication. Thin adapter over agent_registry.authenticate_and_register — the handshake sequence itself is transport-neutral and shared with the long-poll /connect endpoint; this handler only translates the outcome into socket emits/disconnects. Expected data: { "type": "auth", "agent_id": "uuid", "api_key_prefix": "sk_xxx", "signature": "hmac_signature", "timestamp": unix_ms, "nonce": "unique_nonce" (optional but recommended) } """ sid = request.sid ip_address = request.remote_addr logger.info("Auth attempt from agent %s at %s", (data or {}).get('agent_id'), ip_address) session_token, server, error = agent_registry.authenticate_and_register( data or {}, ip_address, user_agent=request.headers.get('User-Agent', ''), transport='ws', socket_id=sid, ) if error == 'rate_limited': emit('auth_response', {'success': False, 'error': 'Rate limit exceeded'}, room=request.sid, namespace='/agent') return if error is not None: messages = { 'missing_fields': 'Missing required fields', 'auth_failed': 'Authentication failed', 'ip_not_allowed': 'IP address not allowed', 'registration_failed': 'Registration failed', } emit('auth_fail', { 'type': 'auth_fail', 'error': messages.get(error, 'Authentication failed') }) disconnect() return # Calculate token expiry (1 hour) expires = int((time.time() + 3600) * 1000) emit('auth_ok', { 'type': 'auth_ok', 'session_token': session_token, 'expires': expires, 'server_id': server.id }) def on_heartbeat(self, data): """ Handle agent heartbeat. Expected data: { "type": "heartbeat", "metrics": { "cpu_percent": float, "memory_percent": float, "disk_percent": float, "container_count": int, "container_running": int } } """ sid = request.sid agent = agent_registry.get_agent_by_socket(sid) if not agent: emit('error', { 'type': 'error', 'code': 'NOT_AUTHENTICATED', 'message': 'Not authenticated' }) return metrics = data.get('metrics', {}) agent_registry.ingest_agent_state(agent, metrics=metrics) emit('heartbeat_ack', {'type': 'heartbeat_ack'}) def on_command_result(self, data): """ Handle command result from agent. Expected data: { "type": "command_result", "command_id": "uuid", "success": bool, "data": any, "error": string, "duration": int (ms) } """ sid = request.sid agent_registry.handle_command_result(sid, data) def on_system_info(self, data): """ Handle system info update from agent. Expected data: { "type": "system_info", "info": { "hostname": string, "os": string, "os_version": string, "platform": string, "architecture": string, "cpu_cores": int, "cpu_model": string, "total_memory": int, "total_disk": int, "docker_version": string, "agent_version": string } } """ sid = request.sid agent = agent_registry.get_agent_by_socket(sid) if not agent: emit('error', { 'type': 'error', 'code': 'NOT_AUTHENTICATED', 'message': 'Not authenticated' }) return info = data.get('info', {}) agent_registry.ingest_agent_state(agent, sysinfo=info) def on_capabilities(self, data): """ Handle agent capability advertisement. Sent by the agent right after auth succeeds (and on every reconnect). The payload tells the panel which feature surfaces this agent can drive — cron, docker, systemd, etc. — and is used to filter target pickers in the UI. Expected data: { "type": "capabilities", "capabilities": {"docker": true, "cron": true, ...}, "platform": "linux", "distro": "ubuntu", "distro_version": "22.04" } """ sid = request.sid agent = agent_registry.get_agent_by_socket(sid) if not agent: emit('error', { 'type': 'error', 'code': 'NOT_AUTHENTICATED', 'message': 'Not authenticated' }) return # Capability side-effects (incl. the WireGuard tunnel reconcile) # happen inside the shared ingest so both transports — this WS # handler and the long-poll /poll body — behave identically. agent_registry.ingest_agent_state(agent, capabilities=data or {}) def on_stream(self, data): """ Handle streaming data from agent (logs, metrics). Expected data: { "type": "stream", "channel": string, "data": any } """ sid = request.sid agent = agent_registry.get_agent_by_socket(sid) if not agent: return channel = data.get('channel') stream_data = data.get('data') # A malformed/partial stream frame — exactly what the poll fallback # exists to tolerate — can omit or null the channel. Calling # .startswith on a non-string would raise AttributeError; drop the # frame instead of letting the handler crash. if not isinstance(channel, str) or not channel: return # Broadcast to subscribers. The channel format determines the room # (grammar owned by app/sockets_rooms.py): # - "metrics" -> server_metrics_room # - "container:xxx:logs" -> server_container_logs_room if channel == 'metrics': room = rooms.server_metrics_room(agent.server_id) elif channel.startswith('container:') and channel.endswith(':logs'): parts = channel.split(':') if len(parts) < 3: return container_id = parts[1] room = rooms.server_container_logs_room(agent.server_id, container_id) else: room = rooms.server_channel_room(agent.server_id, channel) # Emit to the main namespace for UI clients from app import get_socketio socketio = get_socketio() if socketio: socketio.emit( 'server_stream', { 'server_id': agent.server_id, 'channel': channel, 'data': stream_data }, room=room ) def on_error(self, data): """Handle error message from agent""" sid = request.sid agent = agent_registry.get_agent_by_socket(sid) if agent: logger.warning("Error from server %s: %s", agent.server_id, data) def on_credential_update_ack(self, data): """ Handle credential update acknowledgment from agent. Expected data: { "type": "credential_update_ack", "rotation_id": "uuid", "success": bool, "error": string (optional) } """ sid = request.sid agent = agent_registry.get_agent_by_socket(sid) if not agent: emit('error', { 'type': 'error', 'code': 'NOT_AUTHENTICATED', 'message': 'Not authenticated' }) return rotation_id = data.get('rotation_id') success = data.get('success', False) error = data.get('error') logger.info("Credential update ack from %s: success=%s", agent.server_id, success) if not rotation_id: emit('error', { 'type': 'error', 'code': 'MISSING_ROTATION_ID', 'message': 'Missing rotation_id' }) return try: from app import db server = Server.query.get(agent.server_id) if not server: return if success: # Complete the rotation if server.complete_key_rotation(rotation_id): db.session.commit() logger.info("Key rotation completed for server %s", server.id) # Clear nonces for this server since we have new credentials from app.services.nonce_service import nonce_service nonce_service.clear_server_nonces(server.id) else: logger.warning("Key rotation completion failed for server %s", server.id) else: # Agent failed to update credentials, cancel rotation server.cancel_key_rotation() db.session.commit() logger.info("Key rotation cancelled for server %s: %s", server.id, error) except Exception: # Roll back the poisoned session so the next DB user on this # worker doesn't inherit a broken transaction. from app import db db.session.rollback() logger.exception("Error handling credential update ack") def init_agent_gateway(socketio): """Initialize the agent gateway namespace""" agent_registry.init_socketio(socketio) socketio.on_namespace(AgentNamespace('/agent')) logger.info("Agent gateway initialized on /agent namespace")