implement heartbeat command

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}"

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",
+]

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__":

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()

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

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,
+}