diff --git a/app/api/routers/monitor.py b/app/api/routers/monitor.py index b9d7be1..50c5c17 100644 --- a/app/api/routers/monitor.py +++ b/app/api/routers/monitor.py @@ -1,9 +1,10 @@ from datetime import datetime, timedelta, timezone import json +import secrets import uuid from typing import List, Dict -from fastapi import APIRouter, Depends, HTTPException, status, WebSocket, WebSocketDisconnect +from fastapi import APIRouter, Depends, Header, HTTPException, status, WebSocket, WebSocketDisconnect from pydantic import BaseModel from sqlalchemy.orm import Session @@ -171,6 +172,30 @@ def delete_server(server_id: int, db: Session = Depends(get_db), _: models.User return None +@router.post('/admin/servers/{server_id}/api-key') +def generate_api_key(server_id: int, db: Session = Depends(get_db), _: models.User = Depends(require_admin)): + """Generate or regenerate API Key for a server (heartbeat v2)""" + server = db.query(MonitoredServer).filter(MonitoredServer.id == server_id).first() + if not server: + raise HTTPException(status_code=404, detail='Server not found') + # Generate new API key (32 bytes = ~43 chars base64) + api_key = secrets.token_urlsafe(32) + server.api_key = api_key + db.commit() + return {'server_id': server.id, 'api_key': api_key, 'message': 'Store this key securely - it will not be shown again'} + + +@router.delete('/admin/servers/{server_id}/api-key', status_code=status.HTTP_204_NO_CONTENT) +def revoke_api_key(server_id: int, db: Session = Depends(get_db), _: models.User = Depends(require_admin)): + """Revoke API Key for a server""" + server = db.query(MonitoredServer).filter(MonitoredServer.id == server_id).first() + if not server: + raise HTTPException(status_code=404, detail='Server not found') + server.api_key = None + db.commit() + return None + + class ServerHeartbeat(BaseModel): identifier: str openclaw_version: str | None = None @@ -201,6 +226,49 @@ def server_heartbeat(payload: ServerHeartbeat, db: Session = Depends(get_db)): return {'ok': True, 'server_id': server.id, 'last_seen_at': st.last_seen_at} +# Heartbeat v2 with API Key authentication +class TelemetryPayload(BaseModel): + identifier: str + openclaw_version: str | None = None + agents: List[dict] = [] + cpu_pct: float | None = None + mem_pct: float | None = None + disk_pct: float | None = None + swap_pct: float | None = None + load_avg: list[float] | None = None + uptime_seconds: int | None = None + + +@router.post('/server/heartbeat-v2') +def server_heartbeat_v2( + payload: TelemetryPayload, + x_api_key: str = Header(..., alias='X-API-Key', description='API Key from /admin/servers/{id}/api-key'), + db: Session = Depends(get_db) +): + """Server heartbeat using API Key authentication (no challenge_uuid required)""" + # Validate API key + server = db.query(MonitoredServer).filter( + MonitoredServer.api_key == x_api_key, + MonitoredServer.is_enabled == True + ).first() + if not server: + raise HTTPException(status_code=401, detail='Invalid or missing API Key') + # Update server state + st = db.query(ServerState).filter(ServerState.server_id == server.id).first() + if not st: + st = ServerState(server_id=server.id) + db.add(st) + st.openclaw_version = payload.openclaw_version + st.agents_json = json.dumps(payload.agents, ensure_ascii=False) + st.cpu_pct = payload.cpu_pct + st.mem_pct = payload.mem_pct + st.disk_pct = payload.disk_pct + st.swap_pct = payload.swap_pct + st.last_seen_at = datetime.now(timezone.utc) + db.commit() + return {'ok': True, 'server_id': server.id, 'identifier': server.identifier, 'last_seen_at': st.last_seen_at} + + @router.websocket('/server/ws') async def server_ws(websocket: WebSocket): await websocket.accept() diff --git a/app/main.py b/app/main.py index 594531f..eaf1e9b 100644 --- a/app/main.py +++ b/app/main.py @@ -215,6 +215,11 @@ def _migrate_schema(): "DEFAULT 'open'" )) + # --- monitored_servers.api_key for heartbeat v2 --- + if _has_table(db, "monitored_servers") and not _has_column(db, "monitored_servers", "api_key"): + db.execute(text("ALTER TABLE monitored_servers ADD COLUMN api_key VARCHAR(64) NULL")) + db.execute(text("CREATE UNIQUE INDEX idx_monitored_servers_api_key ON monitored_servers (api_key)")) + db.commit() except Exception as e: db.rollback() diff --git a/app/models/monitor.py b/app/models/monitor.py index 533dbcb..21d0fc1 100644 --- a/app/models/monitor.py +++ b/app/models/monitor.py @@ -39,6 +39,7 @@ class MonitoredServer(Base): identifier = Column(String(128), nullable=False, unique=True) display_name = Column(String(128), nullable=True) is_enabled = Column(Boolean, default=True) + api_key = Column(String(64), nullable=True, unique=True, index=True) # API Key for server heartbeat v2 created_by = Column(Integer, nullable=True) created_at = Column(DateTime(timezone=True), server_default=func.now()) diff --git a/tests/test_monitor.py b/tests/test_monitor.py new file mode 100644 index 0000000..5fcc1b8 --- /dev/null +++ b/tests/test_monitor.py @@ -0,0 +1,100 @@ +"""Tests for Monitor API Key authentication and heartbeat v2""" + +import pytest +from unittest.mock import MagicMock, patch +from fastapi.testclient import TestClient + + +# Test helper functions +def test_api_key_generation(): + """Test that API key is generated with correct format""" + import secrets + api_key = secrets.token_urlsafe(32) + # Should be around 43 chars (base64 encoded 32 bytes) + assert len(api_key) >= 40 + assert len(api_key) <= 50 + + +def test_monitored_server_model_has_api_key(): + """Test that MonitoredServer model has api_key field""" + from app.models.monitor import MonitoredServer + # Check the model has the api_key column defined + columns = [c.name for c in MonitoredServer.__table__.columns] + assert 'api_key' in columns + + +def test_api_key_endpoint_exists(): + """Test that /admin/servers/{id}/api-key endpoint is defined""" + from app.api.routers.monitor import router + # Check endpoint is in routes + paths = [r.path for r in router.routes] + assert any('/admin/servers/{server_id}/api-key' in p for p in paths) + + +def test_heartbeat_v2_endpoint_exists(): + """Test that /server/heartbeat-v2 endpoint is defined""" + from app.api.routers.monitor import router + paths = [r.path for r in router.routes] + assert any('/server/heartbeat-v2' in p for p in paths) + + +def test_telemetry_payload_model(): + """Test TelemetryPayload model fields""" + from app.api.routers.monitor import TelemetryPayload + payload = TelemetryPayload( + identifier='test-server', + openclaw_version='1.0.0', + agents=[{'name': 'test', 'status': 'active'}], + cpu_pct=50.0, + mem_pct=60.0, + disk_pct=70.0, + load_avg=[1.0, 2.0, 3.0], + uptime_seconds=3600 + ) + assert payload.identifier == 'test-server' + assert payload.openclaw_version == '1.0.0' + assert len(payload.agents) == 1 + assert payload.load_avg == [1.0, 2.0, 3.0] + + +def test_heartbeat_v2_requires_api_key_header(): + """Test that heartbeat-v2 requires X-API-Key header""" + from app.api.routers.monitor import server_heartbeat_v2 + # The function signature should require x_api_key parameter + import inspect + sig = inspect.signature(server_heartbeat_v2) + params = list(sig.parameters.keys()) + assert 'x_api_key' in params + + +# Integration test placeholders (require running database) +class TestAPIKeyIntegration: + """Integration tests - require database and running server""" + + @pytest.mark.skip(reason="Requires running server") + def test_generate_api_key(self): + """Generate API key for a server""" + # POST /admin/servers/{id}/api-key + # Should return api_key + pass + + @pytest.mark.skip(reason="Requires running server") + def test_heartbeat_v2_with_valid_key(self): + """Send heartbeat with valid API key""" + # POST /server/heartbeat-v2 with X-API-Key header + # Should return ok: true + pass + + @pytest.mark.skip(reason="Requires running server") + def test_heartbeat_v2_with_invalid_key(self): + """Send heartbeat with invalid API key""" + # POST /server/heartbeat-v2 with invalid X-API-Key + # Should return 401 + pass + + @pytest.mark.skip(reason="Requires running server") + def test_revoke_api_key(self): + """Revoke API key""" + # DELETE /admin/servers/{id}/api-key + # Should return 204 + pass \ No newline at end of file