feat: add NetBox plugin store

netbox_plugin/netbox_plugin_store/agent.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import json
+import socket
+import time
+from pathlib import Path
+from typing import Any
+from uuid import UUID, uuid4
+
+from .redaction import redact_text
+
+
+class AgentError(RuntimeError):
+    pass
+
+
+class AgentClient:
+    PROTOCOL_VERSION = 1
+    MAX_MESSAGE_BYTES = 64 * 1024
+
+    def __init__(self, socket_path: Path, *, timeout: int = 30):
+        self.socket_path = Path(socket_path)
+        self.timeout = timeout
+
+    def request(self, method: str, path: str, **values: Any) -> dict[str, Any]:
+        request = {
+            "protocol_version": self.PROTOCOL_VERSION,
+            "method": method,
+            "path": path,
+            **values,
+        }
+        encoded = json.dumps(request, separators=(",", ":"), ensure_ascii=True).encode("utf-8") + b"\n"
+        if len(encoded) > self.MAX_MESSAGE_BYTES:
+            raise AgentError("Agent request exceeds 64 KiB.")
+        if not hasattr(socket, "AF_UNIX"):
+            raise AgentError("Unix sockets are unavailable on this platform.")
+        try:
+            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
+                client.settimeout(self.timeout)
+                client.connect(str(self.socket_path))
+                client.sendall(encoded)
+                client.shutdown(socket.SHUT_WR)
+                response = bytearray()
+                while b"\n" not in response:
+                    chunk = client.recv(4096)
+                    if not chunk:
+                        break
+                    response.extend(chunk)
+                    if len(response) > self.MAX_MESSAGE_BYTES:
+                        raise AgentError("Agent response exceeds 64 KiB.")
+        except (OSError, TimeoutError) as exc:
+            raise AgentError(f"Agent connection failed: {type(exc).__name__}") from exc
+        line, separator, remainder = bytes(response).partition(b"\n")
+        if not separator or remainder:
+            raise AgentError("Agent must return exactly one JSON line.")
+        try:
+            payload = json.loads(line.decode("utf-8"))
+        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+            raise AgentError("Agent returned invalid JSON.") from exc
+        if not isinstance(payload, dict) or payload.get("protocol_version") != self.PROTOCOL_VERSION:
+            raise AgentError("Agent returned an incompatible protocol response.")
+        status = payload.get("status")
+        body = payload.get("body")
+        if not isinstance(status, int) or not isinstance(body, dict):
+            raise AgentError("Agent returned a malformed response.")
+        if status < 200 or status >= 300:
+            message = body.get("error") or body.get("detail") or f"Agent returned status {status}."
+            raise AgentError(redact_text(message, limit=2_000))
+        return body
+
+    def capabilities(self) -> dict[str, Any]:
+        return self.request("GET", "/v1/capabilities")
+
+    def submit_operation(
+        self,
+        *,
+        request_id: str,
+        action: str,
+        plugin_slug: str,
+        version: str | None,
+        approved_payload_sha256: str | None,
+        requested_by: str,
+    ) -> UUID:
+        idempotency_key = uuid4()
+        body = self.request(
+            "POST",
+            "/v1/operations",
+            idempotency_key=str(idempotency_key),
+            body={
+                "request_id": request_id,
+                "action": action,
+                "plugin_slug": plugin_slug,
+                "version": version,
+                "approved_payload_sha256": approved_payload_sha256,
+                "requested_by": requested_by,
+            },
+        )
+        raw_id = body.get("operation_id") or body.get("id")
+        try:
+            return UUID(str(raw_id))
+        except (ValueError, TypeError) as exc:
+            raise AgentError("Agent did not return a valid operation UUID.") from exc
+
+    def wait_for_operation(
+        self,
+        operation_id: UUID,
+        *,
+        overall_timeout: int,
+        poll_interval: float,
+    ) -> dict[str, Any]:
+        deadline = time.monotonic() + overall_timeout
+        while True:
+            body = self.get_operation(operation_id)
+            state = body.get("state") or body.get("status")
+            if state in {"succeeded", "completed"}:
+                result = body.get("result", body)
+                if not isinstance(result, dict):
+                    raise AgentError("Agent operation result is malformed.")
+                return result
+            if state in {"failed", "errored", "cancelled"}:
+                raise AgentError(redact_text(body.get("error") or f"Agent operation {state}."))
+            if state not in {"queued", "pending", "running", "accepted"}:
+                raise AgentError("Agent returned an unknown operation state.")
+            if time.monotonic() >= deadline:
+                raise AgentError("Timed out waiting for the host agent operation.")
+            time.sleep(poll_interval)
+
+    def get_operation(self, operation_id: UUID) -> dict[str, Any]:
+        return self.request("GET", f"/v1/operations/{operation_id}")