66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
"""Agent credential persistence with restrictive local permissions."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AgentState:
|
|
server: str
|
|
agent_id: str
|
|
host_id: str
|
|
token: str
|
|
heartbeat_interval_s: float = 5.0
|
|
|
|
@classmethod
|
|
def from_payload(
|
|
cls,
|
|
server: str,
|
|
payload: Dict[str, Any],
|
|
) -> "AgentState":
|
|
return cls(
|
|
server=server.rstrip("/"),
|
|
agent_id=str(payload["agent_id"]),
|
|
host_id=str(payload["host_id"]),
|
|
token=str(payload["token"]),
|
|
heartbeat_interval_s=float(payload.get("heartbeat_interval_s", 5.0)),
|
|
)
|
|
|
|
|
|
class StateStore:
|
|
def __init__(self, path: Path) -> None:
|
|
self.path = path
|
|
|
|
def load(self) -> Optional[AgentState]:
|
|
try:
|
|
raw = json.loads(self.path.read_text(encoding="utf-8"))
|
|
return AgentState(**raw)
|
|
except FileNotFoundError:
|
|
return None
|
|
except (OSError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
|
raise RuntimeError("Agent state 文件无效: {0}".format(self.path)) from exc
|
|
|
|
def save(self, state: AgentState) -> None:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = self.path.with_suffix(self.path.suffix + ".tmp")
|
|
descriptor = os.open(
|
|
str(temporary),
|
|
os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
|
|
0o600,
|
|
)
|
|
try:
|
|
with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
|
|
json.dump(asdict(state), handle, ensure_ascii=False, indent=2)
|
|
handle.write("\n")
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.chmod(temporary, 0o600)
|
|
os.replace(temporary, self.path)
|
|
finally:
|
|
if temporary.exists():
|
|
temporary.unlink()
|