From f9aed123fd1d1fca559edaae22e6dbcf7c954fcc Mon Sep 17 00:00:00 2001 From: ars Date: Sun, 24 May 2026 18:36:55 +0300 Subject: [PATCH] implement heartbeat command --- pyproject.toml | 1 + src/nexus_sync/client/execute.py | 143 +++++++++++++++++++++--- src/nexus_sync/server/__init__.py | 11 ++ src/nexus_sync/server/__main__.py | 7 +- src/nexus_sync/server/app.py | 84 +++++++++++++++ src/nexus_sync/server/config.py | 16 +++ src/nexus_sync/server/store.py | 143 ++++++++++++++++++++++++ tests/test_execute.py | 100 +++++++++++++++++ tests/test_main.py | 7 -- tests/test_models.py | 2 +- tests/test_server.py | 174 ++++++++++++++++++++++++++++++ 11 files changed, 663 insertions(+), 25 deletions(-) create mode 100644 src/nexus_sync/server/app.py create mode 100644 src/nexus_sync/server/config.py create mode 100644 src/nexus_sync/server/store.py create mode 100644 tests/test_execute.py delete mode 100644 tests/test_main.py create mode 100644 tests/test_server.py diff --git a/pyproject.toml b/pyproject.toml index 7fd83f0..b7c7de9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ [project.optional-dependencies] dev = [ "pytest>=8.0", + "httpx>=0.28.0", "black>=24.0", "pyinstaller>=6.0", "mypy>=1.0", diff --git a/src/nexus_sync/client/execute.py b/src/nexus_sync/client/execute.py index 2862588..b09c601 100644 --- a/src/nexus_sync/client/execute.py +++ b/src/nexus_sync/client/execute.py @@ -1,23 +1,134 @@ -import shlex +import platform import subprocess -from typing import Dict, Optional +from datetime import UTC, datetime +from typing import Any, Callable, Mapping, Sequence + +from nexus_sync.common import Command, CommandKind, CommandResult, CommandResultStatus + +PresetBuilder = Callable[[Mapping[str, Any]], Sequence[str]] + +DEFAULT_OUTPUT_LIMIT_BYTES = 64 * 1024 -def run(command: str, stdin: Optional[str] = None) -> Dict[str, str]: - process = _execute(command, stdin) - - return { - "stdout": process.stdout, - "stderr": process.stderr, - } +def _reject(command: Command, message: str) -> CommandResult: + now = datetime.now(UTC) + return CommandResult( + command_id=command.id, + status=CommandResultStatus.REJECTED, + started_at=now, + finished_at=now, + stderr=message, + ) -def _parse_command(command: str) -> list[str]: - # return command.split() - return shlex.split(command) +def _network_interfaces(args: Mapping[str, Any]) -> Sequence[str]: + if args: + raise ValueError("network_interfaces does not accept arguments") + + system = platform.system().lower() + if system == "windows": + return ["ipconfig"] + if system == "linux": + return ["ip", "addr"] + if system == "darwin": + return ["ifconfig"] + raise ValueError(f"unsupported platform for network_interfaces: {system or 'unknown'}") -def _execute(command: str, stdin: Optional[str] = None) -> subprocess.CompletedProcess: - parsed_command = _parse_command(command) - result = subprocess.run(parsed_command, input=stdin, text=True, capture_output=True, shell=True) - return result +def _hostname(args: Mapping[str, Any]) -> Sequence[str]: + if args: + raise ValueError("hostname does not accept arguments") + return ["hostname"] + + +DEFAULT_PRESETS: dict[str, PresetBuilder] = { + "hostname": _hostname, + "network_interfaces": _network_interfaces, +} + + +def execute_command( + command: Command, + *, + stdin: str | None = None, + presets: Mapping[str, PresetBuilder] = DEFAULT_PRESETS, + output_limit_bytes: int = DEFAULT_OUTPUT_LIMIT_BYTES, +) -> CommandResult: + if command.kind != CommandKind.EXEC: + return _reject(command, f"unsupported command kind: {command.kind}") + + builder = presets.get(command.name) + if builder is None: + return _reject(command, f"unknown command preset: {command.name}") + + try: + argv = list(builder(command.args)) + except ValueError as error: + return _reject(command, str(error)) + + if not argv: + return _reject(command, f"command preset returned empty argv: {command.name}") + + started_at = datetime.now(UTC) + try: + process = subprocess.run( + argv, + input=stdin, + text=True, + capture_output=True, + timeout=command.timeout_seconds, + shell=False, + check=False, + ) + except subprocess.TimeoutExpired as error: + finished_at = datetime.now(UTC) + return CommandResult( + command_id=command.id, + status=CommandResultStatus.TIMED_OUT, + started_at=started_at, + finished_at=finished_at, + stdout=_limit_output(error.stdout or "", output_limit_bytes), + stderr=_limit_output(error.stderr or "", output_limit_bytes), + ) + except OSError as error: + finished_at = datetime.now(UTC) + return CommandResult( + command_id=command.id, + status=CommandResultStatus.FAILED, + started_at=started_at, + finished_at=finished_at, + return_code=None, + stderr=str(error), + ) + + finished_at = datetime.now(UTC) + status = ( + CommandResultStatus.SUCCEEDED if process.returncode == 0 else CommandResultStatus.FAILED + ) + return CommandResult( + command_id=command.id, + status=status, + started_at=started_at, + finished_at=finished_at, + return_code=process.returncode, + stdout=_limit_output(process.stdout, output_limit_bytes), + stderr=_limit_output(process.stderr, output_limit_bytes), + ) + + +def _limit_output(value: str | bytes, limit_bytes: int) -> str: + if isinstance(value, bytes): + value = value.decode(errors="replace") + + encoded = value.encode() + if len(encoded) <= limit_bytes: + return value + + marker = "\n[truncated]" + marker_bytes = marker.encode() + if limit_bytes <= len(marker_bytes): + return encoded[:limit_bytes].decode(errors="ignore") + + content_limit = max(0, limit_bytes - len(marker_bytes)) + truncated = encoded[:content_limit].decode(errors="ignore") + return f"{truncated}{marker}" diff --git a/src/nexus_sync/server/__init__.py b/src/nexus_sync/server/__init__.py index e69de29..dea5d62 100644 --- a/src/nexus_sync/server/__init__.py +++ b/src/nexus_sync/server/__init__.py @@ -0,0 +1,11 @@ +from nexus_sync.server.app import create_app +from nexus_sync.server.config import DEFAULT_COMMAND_POLL_SECONDS, DEFAULT_IDLE_POLL_SECONDS +from nexus_sync.server.store import InMemoryStore, Store + +__all__ = [ + "DEFAULT_COMMAND_POLL_SECONDS", + "DEFAULT_IDLE_POLL_SECONDS", + "InMemoryStore", + "Store", + "create_app", +] diff --git a/src/nexus_sync/server/__main__.py b/src/nexus_sync/server/__main__.py index 7b182bd..baabf72 100644 --- a/src/nexus_sync/server/__main__.py +++ b/src/nexus_sync/server/__main__.py @@ -1,5 +1,10 @@ +import uvicorn + +from nexus_sync.server.app import app + + def main() -> None: - print("nexus-sync server") + uvicorn.run(app, host="0.0.0.0", port=8000) if __name__ == "__main__": diff --git a/src/nexus_sync/server/app.py b/src/nexus_sync/server/app.py new file mode 100644 index 0000000..b1a24bb --- /dev/null +++ b/src/nexus_sync/server/app.py @@ -0,0 +1,84 @@ +from datetime import UTC, datetime + +from fastapi import Depends, FastAPI, Header, HTTPException, Request, status +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse + +from nexus_sync.common import HeartbeatRequest, HeartbeatResponse +from nexus_sync.server.config import ( + DEFAULT_COMMAND_POLL_SECONDS, + DEFAULT_IDLE_POLL_SECONDS, + load_client_tokens, +) +from nexus_sync.server.store import InMemoryStore, Store + + +def create_app( + *, + store: Store | None = None, + client_tokens: dict[str, str] | None = None, +) -> FastAPI: + app = FastAPI(title="nexus-sync", version="0.1.0") + app.state.store = store or InMemoryStore() + app.state.client_tokens = client_tokens if client_tokens is not None else load_client_tokens() + + @app.exception_handler(RequestValidationError) + async def validation_exception_handler( + _request: Request, + error: RequestValidationError, + ) -> JSONResponse: + return JSONResponse( + status_code=status.HTTP_400_BAD_REQUEST, content={"detail": error.errors()} + ) + + def authorize_token(authorization: str | None = Header(default=None)) -> str: + if authorization is None or not authorization.startswith("Bearer "): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="missing bearer token", + ) + + token = authorization.removeprefix("Bearer ").strip() + owner_client_id = _client_id_for_token(app.state.client_tokens, token) + if owner_client_id is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="invalid bearer token", + ) + return owner_client_id + + @app.post("/api/v1/client/heartbeat", response_model=HeartbeatResponse) + def heartbeat( + payload: HeartbeatRequest, + token_client_id: str = Depends(authorize_token), + ) -> HeartbeatResponse: + if token_client_id != payload.client_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="token is not allowed for this client_id", + ) + + now = datetime.now(UTC) + app.state.store.upsert_client(payload, now) + app.state.store.record_command_result(payload, now) + + command = app.state.store.take_next_command(payload.client_id, now) + return HeartbeatResponse( + server_time=now, + next_poll_after_seconds=( + DEFAULT_COMMAND_POLL_SECONDS if command else DEFAULT_IDLE_POLL_SECONDS + ), + command=command, + ) + + return app + + +def _client_id_for_token(client_tokens: dict[str, str], token: str) -> str | None: + for client_id, expected_token in client_tokens.items(): + if token == expected_token: + return client_id + return None + + +app = create_app() diff --git a/src/nexus_sync/server/config.py b/src/nexus_sync/server/config.py new file mode 100644 index 0000000..904dbcf --- /dev/null +++ b/src/nexus_sync/server/config.py @@ -0,0 +1,16 @@ +import os + +DEFAULT_IDLE_POLL_SECONDS = 60 +DEFAULT_COMMAND_POLL_SECONDS = 10 + + +def load_client_tokens() -> dict[str, str]: + raw_tokens = os.environ.get("NEXUS_SYNC_CLIENT_TOKENS", "dev-client:dev-token") + tokens: dict[str, str] = {} + for item in raw_tokens.split(","): + if not item.strip(): + continue + client_id, separator, token = item.partition(":") + if separator and client_id and token: + tokens[client_id] = token + return tokens diff --git a/src/nexus_sync/server/store.py b/src/nexus_sync/server/store.py new file mode 100644 index 0000000..50ada3d --- /dev/null +++ b/src/nexus_sync/server/store.py @@ -0,0 +1,143 @@ +from dataclasses import dataclass, field +from datetime import datetime +from typing import Protocol + +from fastapi import HTTPException, status + +from nexus_sync.common import ( + ClientRecord, + Command, + CommandRecord, + CommandResultRecord, + CommandStatus, + HeartbeatRequest, +) + + +class Store(Protocol): + def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord: + pass + + def record_command_result(self, heartbeat: HeartbeatRequest, now: datetime) -> None: + pass + + def take_next_command(self, client_id: str, now: datetime) -> Command | None: + pass + + +@dataclass +class InMemoryStore: + clients: dict[str, ClientRecord] = field(default_factory=dict) + commands: dict[str, CommandRecord] = field(default_factory=dict) + results: list[CommandResultRecord] = field(default_factory=list) + + def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord: + existing = self.clients.get(heartbeat.client_id) + created_at = existing.created_at if existing else now + token_hash = existing.token_hash if existing else None + + record = ClientRecord( + id=heartbeat.client_id, + hostname=heartbeat.client.hostname, + platform=heartbeat.client.platform, + version=heartbeat.client.version, + created_at=created_at, + last_seen_at=now, + is_active=True, + token_hash=token_hash, + ) + self.clients[heartbeat.client_id] = record + return record + + def enqueue_command(self, command: Command, client_id: str, now: datetime) -> CommandRecord: + record = CommandRecord( + id=command.id, + client_id=client_id, + kind=command.kind, + name=command.name, + args=command.args, + status=CommandStatus.PENDING, + timeout_seconds=command.timeout_seconds, + created_at=now, + ) + self.commands[record.id] = record + return record + + def record_command_result(self, heartbeat: HeartbeatRequest, now: datetime) -> None: + result = heartbeat.last_command_result + if result is None: + return + + command = self.commands.get(result.command_id) + if command is None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="command result references unknown command", + ) + if command.client_id != heartbeat.client_id: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="command result references command for another client", + ) + if command.status in TERMINAL_COMMAND_STATUSES: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="command result references terminal command", + ) + + finished_at = result.finished_at or now + self.commands[command.id] = command.model_copy( + update={ + "status": CommandStatus(result.status.value), + "finished_at": finished_at, + } + ) + self.results.append( + CommandResultRecord( + command_id=result.command_id, + client_id=heartbeat.client_id, + status=result.status, + started_at=result.started_at, + finished_at=finished_at, + return_code=result.return_code, + stdout=result.stdout, + stderr=result.stderr, + received_at=now, + ) + ) + + def take_next_command(self, client_id: str, now: datetime) -> Command | None: + pending = sorted( + ( + command + for command in self.commands.values() + if command.client_id == client_id and command.status == CommandStatus.PENDING + ), + key=lambda command: command.created_at, + ) + if not pending: + return None + + record = pending[0] + self.commands[record.id] = record.model_copy( + update={ + "status": CommandStatus.DELIVERED, + "attempts": record.attempts + 1, + "delivered_at": now, + } + ) + return Command( + id=record.id, + kind=record.kind, + name=record.name, + args=record.args, + timeout_seconds=record.timeout_seconds, + ) + + +TERMINAL_COMMAND_STATUSES = { + CommandStatus.SUCCEEDED, + CommandStatus.FAILED, + CommandStatus.TIMED_OUT, + CommandStatus.REJECTED, +} diff --git a/tests/test_execute.py b/tests/test_execute.py new file mode 100644 index 0000000..128b3e8 --- /dev/null +++ b/tests/test_execute.py @@ -0,0 +1,100 @@ +import subprocess + +from nexus_sync.client.execute import execute_command +from nexus_sync.common import Command, CommandKind, CommandResultStatus + + +def _command(name: str = "hostname", timeout_seconds: int = 30) -> Command: + return Command( + id="cmd_01JY3H8V8W8P3FXDR3S2BM7M6B", + kind=CommandKind.EXEC, + name=name, + args={}, + timeout_seconds=timeout_seconds, + ) + + +def test_execute_command_runs_known_preset_without_shell(monkeypatch) -> None: + calls = [] + + def fake_run(argv, **kwargs): + calls.append((argv, kwargs)) + return subprocess.CompletedProcess(argv, 0, stdout="host\n", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + + result = execute_command(_command()) + + assert result.status == CommandResultStatus.SUCCEEDED + assert result.return_code == 0 + assert result.stdout == "host\n" + assert calls == [ + ( + ["hostname"], + { + "input": None, + "text": True, + "capture_output": True, + "timeout": 30, + "shell": False, + "check": False, + }, + ) + ] + + +def test_execute_command_rejects_unknown_preset() -> None: + result = execute_command(_command(name="rm_everything")) + + assert result.status == CommandResultStatus.REJECTED + assert result.return_code is None + assert "unknown command preset" in result.stderr + + +def test_execute_command_maps_non_zero_exit_to_failed(monkeypatch) -> None: + def fake_run(argv, **kwargs): + return subprocess.CompletedProcess(argv, 2, stdout="", stderr="failed\n") + + monkeypatch.setattr(subprocess, "run", fake_run) + + result = execute_command(_command()) + + assert result.status == CommandResultStatus.FAILED + assert result.return_code == 2 + assert result.stderr == "failed\n" + + +def test_execute_command_maps_timeout(monkeypatch) -> None: + def fake_run(argv, **kwargs): + raise subprocess.TimeoutExpired(argv, timeout=kwargs["timeout"], output="partial") + + monkeypatch.setattr(subprocess, "run", fake_run) + + result = execute_command(_command(timeout_seconds=1)) + + assert result.status == CommandResultStatus.TIMED_OUT + assert result.return_code is None + assert result.started_at is not None + assert result.finished_at is not None + + +def test_execute_command_truncates_output(monkeypatch) -> None: + def fake_run(argv, **kwargs): + return subprocess.CompletedProcess(argv, 0, stdout="abcdefghijklmnopqrstuvwxyz", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + + result = execute_command(_command(), output_limit_bytes=20) + + assert result.stdout.endswith("\n[truncated]") + assert len(result.stdout.encode()) <= 20 + + +def test_execute_command_rejects_preset_validation_error() -> None: + command = _command() + invalid = command.model_copy(update={"args": {"unexpected": True}}) + + result = execute_command(invalid) + + assert result.status == CommandResultStatus.REJECTED + assert "does not accept arguments" in result.stderr diff --git a/tests/test_main.py b/tests/test_main.py deleted file mode 100644 index c2f4f94..0000000 --- a/tests/test_main.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -foobuzz test -""" - - -def test_foo(): - pass diff --git a/tests/test_models.py b/tests/test_models.py index ff6633a..a8ca9db 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -15,7 +15,7 @@ from nexus_sync.common import ( ) -def test_heartbeat_request_accepts_current_mvp_payload() -> None: +def test_heartbeat_request_accepts_payload() -> None: payload = HeartbeatRequest( client_id="macbook-pro-01", observed_at=datetime(2026, 5, 24, 13, 20, 30, tzinfo=UTC), diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..c5304b7 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,174 @@ +from datetime import UTC, datetime + +from fastapi.testclient import TestClient + +from nexus_sync.common import Command, CommandKind, CommandResultStatus, CommandStatus +from nexus_sync.server import ( + DEFAULT_COMMAND_POLL_SECONDS, + DEFAULT_IDLE_POLL_SECONDS, + Store, + InMemoryStore, + create_app, +) + + +def _heartbeat_payload(client_id: str = "macbook-pro-01", result: dict | None = None) -> dict: + return { + "client_id": client_id, + "observed_at": "2026-05-24T13:20:30Z", + "client": { + "hostname": "macbook-pro.local", + "platform": "darwin", + "version": "0.1.0", + }, + "state": { + "local_time": "2026-05-24T16:20:30+03:00", + "uptime_seconds": 1200, + }, + "last_command_result": result, + } + + +def _client(store: Store | None = None) -> TestClient: + app = create_app( + store=store or InMemoryStore(), + client_tokens={"macbook-pro-01": "client-token", "other-client": "other-token"}, + ) + return TestClient(app) + + +def test_heartbeat_requires_bearer_token() -> None: + client = _client() + + response = client.post("/api/v1/client/heartbeat", json=_heartbeat_payload()) + + assert response.status_code == 401 + + +def test_heartbeat_rejects_token_for_another_client() -> None: + client = _client() + + response = client.post( + "/api/v1/client/heartbeat", + json=_heartbeat_payload(), + headers={"Authorization": "Bearer other-token"}, + ) + + assert response.status_code == 403 + + +def test_heartbeat_returns_bad_request_for_invalid_payload() -> None: + payload = _heartbeat_payload() + payload["unexpected"] = True + client = _client() + + response = client.post( + "/api/v1/client/heartbeat", + json=payload, + headers={"Authorization": "Bearer client-token"}, + ) + + assert response.status_code == 400 + + +def test_heartbeat_accepts_state_without_command() -> None: + store = InMemoryStore() + client = _client(store) + + response = client.post( + "/api/v1/client/heartbeat", + json=_heartbeat_payload(), + headers={"Authorization": "Bearer client-token"}, + ) + + assert response.status_code == 200 + assert response.json()["command"] is None + assert response.json()["next_poll_after_seconds"] == DEFAULT_IDLE_POLL_SECONDS + assert store.clients["macbook-pro-01"].hostname == "macbook-pro.local" + + +def test_heartbeat_delivers_pending_command_and_marks_it_delivered() -> None: + store = InMemoryStore() + store.enqueue_command( + Command( + id="cmd_01JY3H8V8W8P3FXDR3S2BM7M6B", + kind=CommandKind.EXEC, + name="hostname", + args={}, + timeout_seconds=30, + ), + client_id="macbook-pro-01", + now=datetime(2026, 5, 24, 13, 20, 0, tzinfo=UTC), + ) + client = _client(store) + + response = client.post( + "/api/v1/client/heartbeat", + json=_heartbeat_payload(), + headers={"Authorization": "Bearer client-token"}, + ) + + body = response.json() + assert response.status_code == 200 + assert body["next_poll_after_seconds"] == DEFAULT_COMMAND_POLL_SECONDS + assert body["command"]["id"] == "cmd_01JY3H8V8W8P3FXDR3S2BM7M6B" + assert store.commands["cmd_01JY3H8V8W8P3FXDR3S2BM7M6B"].status == CommandStatus.DELIVERED + assert store.commands["cmd_01JY3H8V8W8P3FXDR3S2BM7M6B"].attempts == 1 + + +def test_heartbeat_records_command_result() -> None: + store = InMemoryStore() + store.enqueue_command( + Command( + id="cmd_01JY3H8V8W8P3FXDR3S2BM7M6B", + kind=CommandKind.EXEC, + name="hostname", + args={}, + timeout_seconds=30, + ), + client_id="macbook-pro-01", + now=datetime(2026, 5, 24, 13, 20, 0, tzinfo=UTC), + ) + store.take_next_command("macbook-pro-01", datetime(2026, 5, 24, 13, 20, 1, tzinfo=UTC)) + client = _client(store) + + response = client.post( + "/api/v1/client/heartbeat", + json=_heartbeat_payload( + result={ + "command_id": "cmd_01JY3H8V8W8P3FXDR3S2BM7M6B", + "status": "succeeded", + "started_at": "2026-05-24T13:20:02Z", + "finished_at": "2026-05-24T13:20:03Z", + "return_code": 0, + "stdout": "host\n", + "stderr": "", + } + ), + headers={"Authorization": "Bearer client-token"}, + ) + + assert response.status_code == 200 + assert store.commands["cmd_01JY3H8V8W8P3FXDR3S2BM7M6B"].status == CommandStatus.SUCCEEDED + assert len(store.results) == 1 + assert store.results[0].status == CommandResultStatus.SUCCEEDED + + +def test_heartbeat_rejects_unknown_command_result() -> None: + client = _client() + + response = client.post( + "/api/v1/client/heartbeat", + json=_heartbeat_payload( + result={ + "command_id": "cmd_unknown", + "status": "succeeded", + "return_code": 0, + "stdout": "", + "stderr": "", + } + ), + headers={"Authorization": "Bearer client-token"}, + ) + + assert response.status_code == 409