102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass, field, replace
|
|
from pathlib import Path
|
|
from secrets import token_urlsafe
|
|
from uuid import uuid4
|
|
|
|
|
|
class HostIdentityStateError(RuntimeError):
|
|
"""Raised when persisted Host enrollment identity is missing or invalid."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HostIdentityState:
|
|
agent_instance_id: str
|
|
token: str = field(repr=False)
|
|
host_id: str | None = None
|
|
|
|
|
|
class HostIdentityStore:
|
|
def __init__(self, path: str | Path) -> None:
|
|
self.path = Path(path)
|
|
|
|
def load(self) -> HostIdentityState | None:
|
|
if not self.path.exists():
|
|
return None
|
|
try:
|
|
payload = json.loads(self.path.read_text(encoding="utf-8"))
|
|
except (OSError, ValueError, json.JSONDecodeError) as exc:
|
|
raise HostIdentityStateError("Host identity state is unreadable") from exc
|
|
if not isinstance(payload, dict):
|
|
raise HostIdentityStateError("Host identity state must be an object")
|
|
agent_instance_id = payload.get("agent_instance_id")
|
|
token = payload.get("token")
|
|
host_id = payload.get("host_id")
|
|
if (
|
|
not isinstance(agent_instance_id, str)
|
|
or not agent_instance_id
|
|
or not isinstance(token, str)
|
|
or len(token) < 32
|
|
or (host_id is not None and (not isinstance(host_id, str) or not host_id))
|
|
):
|
|
raise HostIdentityStateError("Host identity state is invalid")
|
|
return HostIdentityState(
|
|
agent_instance_id=agent_instance_id,
|
|
token=token,
|
|
host_id=host_id,
|
|
)
|
|
|
|
def load_or_create(self) -> HostIdentityState:
|
|
existing = self.load()
|
|
if existing is not None:
|
|
return existing
|
|
state = HostIdentityState(
|
|
agent_instance_id=f"agent-{uuid4().hex}",
|
|
token=token_urlsafe(48),
|
|
)
|
|
self._write(state)
|
|
return state
|
|
|
|
def complete(self, state: HostIdentityState, host_id: str) -> HostIdentityState:
|
|
if not host_id:
|
|
raise ValueError("host_id must not be empty")
|
|
current = self.load()
|
|
if current is not None and (
|
|
current.agent_instance_id != state.agent_instance_id
|
|
or current.token != state.token
|
|
):
|
|
raise HostIdentityStateError("Host identity changed during enrollment")
|
|
completed = replace(state, host_id=host_id)
|
|
self._write(completed)
|
|
return completed
|
|
|
|
def _write(self, state: HostIdentityState) -> None:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = self.path.with_name(f".{self.path.name}.{uuid4().hex}.tmp")
|
|
payload = {
|
|
"agent_instance_id": state.agent_instance_id,
|
|
"token": state.token,
|
|
"host_id": state.host_id,
|
|
}
|
|
try:
|
|
temporary.write_text(
|
|
json.dumps(payload, ensure_ascii=True, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
_restrict_permissions(temporary)
|
|
os.replace(temporary, self.path)
|
|
_restrict_permissions(self.path)
|
|
finally:
|
|
if temporary.exists():
|
|
temporary.unlink()
|
|
|
|
|
|
def _restrict_permissions(path: Path) -> None:
|
|
try:
|
|
path.chmod(0o600)
|
|
except OSError:
|
|
return
|