add client heartbeat runtime
This commit is contained in:
ars
@@ -1,6 +1,6 @@
|
||||
def main() -> None:
|
||||
print("nexus-sync client")
|
||||
import sys
|
||||
|
||||
from nexus_sync.client.runtime import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
raise SystemExit(main(sys.argv[1:]))
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import socket
|
||||
import sys
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from typing import Callable, Mapping
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from nexus_sync.client.config import load_command_access_policy
|
||||
from nexus_sync.client.execute import CommandAccessPolicy, execute_command
|
||||
from nexus_sync.common import (
|
||||
ClientInfo,
|
||||
ClientPlatform,
|
||||
ClientState,
|
||||
CommandResult,
|
||||
HeartbeatRequest,
|
||||
HeartbeatResponse,
|
||||
)
|
||||
|
||||
SERVER_URL_ENV = "NEXUS_SYNC_SERVER_URL"
|
||||
CLIENT_ID_ENV = "NEXUS_SYNC_CLIENT_ID"
|
||||
CLIENT_TOKEN_ENV = "NEXUS_SYNC_CLIENT_TOKEN"
|
||||
CLIENT_VERSION = "0.1.0"
|
||||
HEARTBEAT_PATH = "/api/v1/client/heartbeat"
|
||||
|
||||
|
||||
class ClientConfigError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class HeartbeatError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClientConfig:
|
||||
server_url: str
|
||||
client_id: str
|
||||
token: str
|
||||
command_access_policy: CommandAccessPolicy
|
||||
|
||||
|
||||
def load_client_config(env: Mapping[str, str] = os.environ) -> ClientConfig:
|
||||
server_url = _required_env(env, SERVER_URL_ENV).rstrip("/")
|
||||
client_id = _required_env(env, CLIENT_ID_ENV)
|
||||
token = _required_env(env, CLIENT_TOKEN_ENV)
|
||||
return ClientConfig(
|
||||
server_url=server_url,
|
||||
client_id=client_id,
|
||||
token=token,
|
||||
command_access_policy=load_command_access_policy(env),
|
||||
)
|
||||
|
||||
|
||||
def build_heartbeat_request(config: ClientConfig) -> HeartbeatRequest:
|
||||
now = datetime.now(UTC)
|
||||
return HeartbeatRequest(
|
||||
client_id=config.client_id,
|
||||
observed_at=now,
|
||||
client=ClientInfo(
|
||||
hostname=socket.gethostname(),
|
||||
platform=_current_platform(),
|
||||
version=CLIENT_VERSION,
|
||||
),
|
||||
state=ClientState(
|
||||
local_time=datetime.now().astimezone(),
|
||||
uptime_seconds=None,
|
||||
),
|
||||
last_command_result=None,
|
||||
)
|
||||
|
||||
|
||||
def send_heartbeat(
|
||||
config: ClientConfig,
|
||||
heartbeat: HeartbeatRequest,
|
||||
*,
|
||||
opener: Callable[[urllib.request.Request], object] = urllib.request.urlopen,
|
||||
) -> HeartbeatResponse:
|
||||
payload = heartbeat.model_dump_json().encode()
|
||||
request = urllib.request.Request(
|
||||
f"{config.server_url}{HEARTBEAT_PATH}",
|
||||
data=payload,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {config.token}",
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
|
||||
try:
|
||||
with opener(request) as response:
|
||||
body = response.read()
|
||||
except urllib.error.HTTPError as error:
|
||||
detail = error.read().decode(errors="replace")
|
||||
raise HeartbeatError(f"heartbeat failed with HTTP {error.code}: {detail}") from error
|
||||
except urllib.error.URLError as error:
|
||||
raise HeartbeatError(f"heartbeat request failed: {error.reason}") from error
|
||||
except OSError as error:
|
||||
raise HeartbeatError(f"heartbeat request failed: {error}") from error
|
||||
|
||||
try:
|
||||
return HeartbeatResponse.model_validate_json(body)
|
||||
except ValidationError as error:
|
||||
raise HeartbeatError(f"heartbeat response is invalid: {error}") from error
|
||||
|
||||
|
||||
def run_once(
|
||||
config: ClientConfig,
|
||||
*,
|
||||
heartbeat_sender: Callable[
|
||||
[ClientConfig, HeartbeatRequest],
|
||||
HeartbeatResponse,
|
||||
] = send_heartbeat,
|
||||
executor: Callable[..., CommandResult] = execute_command,
|
||||
) -> CommandResult | None:
|
||||
response = heartbeat_sender(config, build_heartbeat_request(config))
|
||||
if response.command is None:
|
||||
return None
|
||||
|
||||
return executor(
|
||||
response.command,
|
||||
access_policy=config.command_access_policy,
|
||||
)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
_ = argv
|
||||
try:
|
||||
config = load_client_config()
|
||||
result = run_once(config)
|
||||
except (ClientConfigError, HeartbeatError, ValueError) as error:
|
||||
print(f"nexus-sync client error: {error}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
if result is None:
|
||||
print("heartbeat accepted; no command")
|
||||
return 0
|
||||
|
||||
print(
|
||||
json.dumps(
|
||||
result.model_dump(mode="json"),
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def _required_env(env: Mapping[str, str], name: str) -> str:
|
||||
value = env.get(name)
|
||||
if value is None or not value.strip():
|
||||
raise ClientConfigError(f"{name} is required")
|
||||
return value.strip()
|
||||
|
||||
|
||||
def _current_platform() -> ClientPlatform:
|
||||
system = platform.system().lower()
|
||||
if system == "linux":
|
||||
return ClientPlatform.LINUX
|
||||
if system == "darwin":
|
||||
return ClientPlatform.DARWIN
|
||||
if system == "windows":
|
||||
return ClientPlatform.WINDOWS
|
||||
return ClientPlatform.UNKNOWN
|
||||
@@ -0,0 +1,237 @@
|
||||
import io
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from nexus_sync.client.execute import CommandAccessPolicy
|
||||
from nexus_sync.client.runtime import (
|
||||
CLIENT_ID_ENV,
|
||||
CLIENT_TOKEN_ENV,
|
||||
SERVER_URL_ENV,
|
||||
ClientConfig,
|
||||
ClientConfigError,
|
||||
HeartbeatError,
|
||||
build_heartbeat_request,
|
||||
load_client_config,
|
||||
main,
|
||||
run_once,
|
||||
send_heartbeat,
|
||||
)
|
||||
from nexus_sync.common import (
|
||||
Command,
|
||||
CommandKind,
|
||||
CommandResult,
|
||||
CommandResultStatus,
|
||||
HeartbeatRequest,
|
||||
HeartbeatResponse,
|
||||
)
|
||||
|
||||
|
||||
def _config(policy: CommandAccessPolicy | None = None) -> ClientConfig:
|
||||
return ClientConfig(
|
||||
server_url="https://nexus.example.test",
|
||||
client_id="macbook-pro-01",
|
||||
token="client-token",
|
||||
command_access_policy=policy or CommandAccessPolicy.deny_all(),
|
||||
)
|
||||
|
||||
|
||||
def test_load_client_config_reads_required_env_and_normalizes_server_url() -> None:
|
||||
config = load_client_config(
|
||||
{
|
||||
SERVER_URL_ENV: "https://nexus.example.test/",
|
||||
CLIENT_ID_ENV: "macbook-pro-01",
|
||||
CLIENT_TOKEN_ENV: "client-token",
|
||||
"NEXUS_SYNC_ALLOWED_COMMANDS": "hostname",
|
||||
}
|
||||
)
|
||||
|
||||
assert config.server_url == "https://nexus.example.test"
|
||||
assert config.client_id == "macbook-pro-01"
|
||||
assert config.token == "client-token"
|
||||
assert config.command_access_policy.allows("hostname")
|
||||
assert not config.command_access_policy.allows("network_interfaces")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("missing_name", [SERVER_URL_ENV, CLIENT_ID_ENV, CLIENT_TOKEN_ENV])
|
||||
def test_load_client_config_requires_env_values(missing_name: str) -> None:
|
||||
env = {
|
||||
SERVER_URL_ENV: "https://nexus.example.test",
|
||||
CLIENT_ID_ENV: "macbook-pro-01",
|
||||
CLIENT_TOKEN_ENV: "client-token",
|
||||
}
|
||||
del env[missing_name]
|
||||
|
||||
with pytest.raises(ClientConfigError, match=missing_name):
|
||||
load_client_config(env)
|
||||
|
||||
|
||||
def test_build_heartbeat_request_contains_client_state(monkeypatch) -> None:
|
||||
monkeypatch.setattr("socket.gethostname", lambda: "macbook-pro.local")
|
||||
monkeypatch.setattr("platform.system", lambda: "Darwin")
|
||||
|
||||
heartbeat = build_heartbeat_request(_config())
|
||||
serialized = heartbeat.model_dump(mode="json")
|
||||
|
||||
assert heartbeat.client_id == "macbook-pro-01"
|
||||
assert heartbeat.client.hostname == "macbook-pro.local"
|
||||
assert heartbeat.client.platform == "darwin"
|
||||
assert heartbeat.client.version == "0.1.0"
|
||||
assert heartbeat.last_command_result is None
|
||||
assert serialized["client_id"] == "macbook-pro-01"
|
||||
assert serialized["state"]["uptime_seconds"] is None
|
||||
|
||||
|
||||
def test_send_heartbeat_posts_json_with_bearer_token() -> None:
|
||||
captured = {}
|
||||
|
||||
def fake_opener(request: urllib.request.Request):
|
||||
captured["url"] = request.full_url
|
||||
captured["authorization"] = request.get_header("Authorization")
|
||||
captured["content_type"] = request.get_header("Content-type")
|
||||
captured["data"] = request.data
|
||||
return _Response(
|
||||
HeartbeatResponse(
|
||||
server_time=datetime(2026, 5, 24, 13, 20, 30, tzinfo=UTC),
|
||||
next_poll_after_seconds=60,
|
||||
command=None,
|
||||
)
|
||||
.model_dump_json()
|
||||
.encode()
|
||||
)
|
||||
|
||||
response = send_heartbeat(
|
||||
_config(),
|
||||
HeartbeatRequest(
|
||||
client_id="macbook-pro-01",
|
||||
observed_at=datetime(2026, 5, 24, 13, 20, 30, tzinfo=UTC),
|
||||
client={
|
||||
"hostname": "macbook-pro.local",
|
||||
"platform": "darwin",
|
||||
"version": "0.1.0",
|
||||
},
|
||||
state={
|
||||
"local_time": datetime(2026, 5, 24, 16, 20, 30, tzinfo=UTC),
|
||||
"uptime_seconds": None,
|
||||
},
|
||||
last_command_result=None,
|
||||
),
|
||||
opener=fake_opener,
|
||||
)
|
||||
|
||||
assert response.command is None
|
||||
assert captured["url"] == "https://nexus.example.test/api/v1/client/heartbeat"
|
||||
assert captured["authorization"] == "Bearer client-token"
|
||||
assert captured["content_type"] == "application/json"
|
||||
assert b"macbook-pro-01" in captured["data"]
|
||||
|
||||
|
||||
def test_send_heartbeat_maps_http_error_to_runtime_error() -> None:
|
||||
def fake_opener(_request: urllib.request.Request):
|
||||
raise urllib.error.HTTPError(
|
||||
url="https://nexus.example.test/api/v1/client/heartbeat",
|
||||
code=401,
|
||||
msg="Unauthorized",
|
||||
hdrs={},
|
||||
fp=io.BytesIO(b'{"detail":"invalid bearer token"}'),
|
||||
)
|
||||
|
||||
with pytest.raises(HeartbeatError, match="HTTP 401"):
|
||||
send_heartbeat(_config(), build_heartbeat_request(_config()), opener=fake_opener)
|
||||
|
||||
|
||||
def test_run_once_without_command_does_not_call_executor() -> None:
|
||||
executor_called = False
|
||||
|
||||
def fake_sender(_config: ClientConfig, _heartbeat: HeartbeatRequest) -> HeartbeatResponse:
|
||||
return HeartbeatResponse(
|
||||
server_time=datetime(2026, 5, 24, 13, 20, 30, tzinfo=UTC),
|
||||
next_poll_after_seconds=60,
|
||||
command=None,
|
||||
)
|
||||
|
||||
def fake_executor(*_args, **_kwargs):
|
||||
nonlocal executor_called
|
||||
executor_called = True
|
||||
raise AssertionError("executor should not be called")
|
||||
|
||||
result = run_once(_config(), heartbeat_sender=fake_sender, executor=fake_executor)
|
||||
|
||||
assert result is None
|
||||
assert executor_called is False
|
||||
|
||||
|
||||
def test_run_once_executes_command_with_configured_access_policy() -> None:
|
||||
policy = CommandAccessPolicy.allow(["hostname"])
|
||||
seen = {}
|
||||
|
||||
def fake_sender(_config: ClientConfig, _heartbeat: HeartbeatRequest) -> HeartbeatResponse:
|
||||
return HeartbeatResponse(
|
||||
server_time=datetime(2026, 5, 24, 13, 20, 30, tzinfo=UTC),
|
||||
next_poll_after_seconds=10,
|
||||
command=Command(
|
||||
id="cmd_01JY3H8V8W8P3FXDR3S2BM7M6B",
|
||||
kind=CommandKind.EXEC,
|
||||
name="hostname",
|
||||
args={},
|
||||
timeout_seconds=30,
|
||||
),
|
||||
)
|
||||
|
||||
def fake_executor(command: Command, **kwargs) -> CommandResult:
|
||||
seen["command"] = command
|
||||
seen["access_policy"] = kwargs["access_policy"]
|
||||
return CommandResult(
|
||||
command_id=command.id,
|
||||
status=CommandResultStatus.SUCCEEDED,
|
||||
return_code=0,
|
||||
stdout="host\n",
|
||||
)
|
||||
|
||||
result = run_once(_config(policy), heartbeat_sender=fake_sender, executor=fake_executor)
|
||||
|
||||
assert result is not None
|
||||
assert result.status == CommandResultStatus.SUCCEEDED
|
||||
assert seen["command"].name == "hostname"
|
||||
assert seen["access_policy"] == policy
|
||||
|
||||
|
||||
def test_main_returns_non_zero_for_missing_config(monkeypatch, capsys) -> None:
|
||||
monkeypatch.delenv(SERVER_URL_ENV, raising=False)
|
||||
monkeypatch.delenv(CLIENT_ID_ENV, raising=False)
|
||||
monkeypatch.delenv(CLIENT_TOKEN_ENV, raising=False)
|
||||
|
||||
exit_code = main([])
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert exit_code == 1
|
||||
assert SERVER_URL_ENV in captured.err
|
||||
|
||||
|
||||
def test_main_prints_success_without_command(monkeypatch, capsys) -> None:
|
||||
monkeypatch.setenv(SERVER_URL_ENV, "https://nexus.example.test")
|
||||
monkeypatch.setenv(CLIENT_ID_ENV, "macbook-pro-01")
|
||||
monkeypatch.setenv(CLIENT_TOKEN_ENV, "client-token")
|
||||
monkeypatch.setattr("nexus_sync.client.runtime.run_once", lambda _config: None)
|
||||
|
||||
exit_code = main([])
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert exit_code == 0
|
||||
assert captured.out == "heartbeat accepted; no command\n"
|
||||
|
||||
|
||||
class _Response:
|
||||
def __init__(self, body: bytes) -> None:
|
||||
self._body = body
|
||||
|
||||
def __enter__(self) -> "_Response":
|
||||
return self
|
||||
|
||||
def __exit__(self, *_args) -> None:
|
||||
return None
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self._body
|
||||
Reference in New Issue
Block a user