mirror of
https://github.com/bunny-lab-io/Borealis.git
synced 2025-10-26 17:41:58 -06:00
149 lines
4.8 KiB
Python
149 lines
4.8 KiB
Python
"""Agent REST endpoints for device communication."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from functools import wraps
|
|
from typing import Any, Callable, Optional, TypeVar, cast
|
|
|
|
from flask import Blueprint, Flask, current_app, g, jsonify, request
|
|
|
|
from Data.Engine.builders.device_auth import DeviceAuthRequestBuilder
|
|
from Data.Engine.domain.device_auth import DeviceAuthContext, DeviceAuthFailure
|
|
from Data.Engine.services.container import EngineServiceContainer
|
|
from Data.Engine.services.devices.device_inventory_service import (
|
|
DeviceDetailsError,
|
|
DeviceHeartbeatError,
|
|
)
|
|
|
|
AGENT_CONTEXT_HEADER = "X-Borealis-Agent-Context"
|
|
|
|
blueprint = Blueprint("engine_agent", __name__)
|
|
|
|
F = TypeVar("F", bound=Callable[..., Any])
|
|
|
|
|
|
def _services() -> EngineServiceContainer:
|
|
return cast(EngineServiceContainer, current_app.extensions["engine_services"])
|
|
|
|
|
|
def require_device_auth(func: F) -> F:
|
|
@wraps(func)
|
|
def wrapper(*args: Any, **kwargs: Any):
|
|
services = _services()
|
|
builder = (
|
|
DeviceAuthRequestBuilder()
|
|
.with_authorization(request.headers.get("Authorization"))
|
|
.with_http_method(request.method)
|
|
.with_htu(request.url)
|
|
.with_service_context(request.headers.get(AGENT_CONTEXT_HEADER))
|
|
.with_dpop_proof(request.headers.get("DPoP"))
|
|
)
|
|
try:
|
|
auth_request = builder.build()
|
|
context = services.device_auth.authenticate(auth_request, path=request.path)
|
|
except DeviceAuthFailure as exc:
|
|
payload = exc.to_dict()
|
|
response = jsonify(payload)
|
|
if exc.retry_after is not None:
|
|
response.headers["Retry-After"] = str(int(math.ceil(exc.retry_after)))
|
|
return response, exc.http_status
|
|
|
|
g.device_auth = context
|
|
try:
|
|
return func(*args, **kwargs)
|
|
finally:
|
|
g.pop("device_auth", None)
|
|
|
|
return cast(F, wrapper)
|
|
|
|
|
|
def register(app: Flask, _services: EngineServiceContainer) -> None:
|
|
if "engine_agent" not in app.blueprints:
|
|
app.register_blueprint(blueprint)
|
|
|
|
|
|
@blueprint.route("/api/agent/heartbeat", methods=["POST"])
|
|
@require_device_auth
|
|
def heartbeat() -> Any:
|
|
services = _services()
|
|
payload = request.get_json(force=True, silent=True) or {}
|
|
context = cast(DeviceAuthContext, g.device_auth)
|
|
|
|
try:
|
|
services.device_inventory.record_heartbeat(context=context, payload=payload)
|
|
except DeviceHeartbeatError as exc:
|
|
error_payload = {"error": exc.code}
|
|
if exc.code == "device_not_registered":
|
|
return jsonify(error_payload), 404
|
|
if exc.code == "storage_conflict":
|
|
return jsonify(error_payload), 409
|
|
current_app.logger.exception(
|
|
"device-heartbeat-error guid=%s code=%s", context.identity.guid.value, exc.code
|
|
)
|
|
return jsonify(error_payload), 500
|
|
|
|
return jsonify({"status": "ok", "poll_after_ms": 15000})
|
|
|
|
|
|
@blueprint.route("/api/agent/script/request", methods=["POST"])
|
|
@require_device_auth
|
|
def script_request() -> Any:
|
|
services = _services()
|
|
context = cast(DeviceAuthContext, g.device_auth)
|
|
|
|
signing_key: Optional[str] = None
|
|
signer = services.script_signer
|
|
if signer is not None:
|
|
try:
|
|
signing_key = signer.public_base64_spki()
|
|
except Exception as exc: # pragma: no cover - defensive logging
|
|
current_app.logger.warning("script-signer-unavailable: %s", exc)
|
|
|
|
status = "quarantined" if context.is_quarantined else "idle"
|
|
poll_after = 60000 if context.is_quarantined else 30000
|
|
|
|
response = {
|
|
"status": status,
|
|
"poll_after_ms": poll_after,
|
|
"sig_alg": "ed25519",
|
|
}
|
|
if signing_key:
|
|
response["signing_key"] = signing_key
|
|
return jsonify(response)
|
|
|
|
|
|
@blueprint.route("/api/agent/details", methods=["POST"])
|
|
@require_device_auth
|
|
def save_details() -> Any:
|
|
services = _services()
|
|
payload = request.get_json(force=True, silent=True) or {}
|
|
context = cast(DeviceAuthContext, g.device_auth)
|
|
|
|
try:
|
|
services.device_inventory.save_agent_details(context=context, payload=payload)
|
|
except DeviceDetailsError as exc:
|
|
error_payload = {"error": exc.code}
|
|
if exc.code == "invalid_payload":
|
|
return jsonify(error_payload), 400
|
|
if exc.code in {"fingerprint_mismatch", "guid_mismatch"}:
|
|
return jsonify(error_payload), 403
|
|
if exc.code == "device_not_registered":
|
|
return jsonify(error_payload), 404
|
|
current_app.logger.exception(
|
|
"device-details-error guid=%s code=%s", context.identity.guid.value, exc.code
|
|
)
|
|
return jsonify(error_payload), 500
|
|
|
|
return jsonify({"status": "ok"})
|
|
|
|
|
|
__all__ = [
|
|
"register",
|
|
"blueprint",
|
|
"heartbeat",
|
|
"script_request",
|
|
"save_details",
|
|
"require_device_auth",
|
|
]
|