mirror of
https://github.com/bunny-lab-io/Borealis.git
synced 2025-12-15 23:25:48 -07:00
405 lines
15 KiB
Python
405 lines
15 KiB
Python
# ======================================================
|
|
# Data\Engine\services\WebSocket\__init__.py
|
|
# Description: Socket.IO handlers for Engine runtime quick job updates and realtime notifications.
|
|
#
|
|
# API Endpoints (if applicable): None
|
|
# ======================================================
|
|
|
|
"""WebSocket service registration for the Borealis Engine runtime."""
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import sqlite3
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, Optional
|
|
|
|
from flask_socketio import SocketIO
|
|
from flask import session, request
|
|
|
|
from ...database import initialise_engine_database
|
|
from ...server import EngineContext
|
|
from .Agent.ReverseTunnel import (
|
|
ReverseTunnelService,
|
|
TunnelBridge,
|
|
decode_frame,
|
|
TunnelFrame,
|
|
)
|
|
from ..API import _make_db_conn_factory, _make_service_logger
|
|
|
|
|
|
def _now_ts() -> int:
|
|
return int(time.time())
|
|
|
|
|
|
def _normalize_text(value: Any) -> str:
|
|
if value is None:
|
|
return ""
|
|
if isinstance(value, bytes):
|
|
try:
|
|
return value.decode("utf-8")
|
|
except Exception:
|
|
return value.decode("utf-8", errors="replace")
|
|
return str(value)
|
|
|
|
|
|
@dataclass
|
|
class EngineRealtimeAdapters:
|
|
context: EngineContext
|
|
db_conn_factory: Callable[[], sqlite3.Connection] = field(init=False)
|
|
service_log: Callable[[str, str, Optional[str]], None] = field(init=False)
|
|
|
|
def __post_init__(self) -> None:
|
|
initialise_engine_database(self.context.database_path, logger=self.context.logger)
|
|
self.db_conn_factory = _make_db_conn_factory(self.context.database_path)
|
|
|
|
log_file = str(
|
|
self.context.config.get("log_file")
|
|
or self.context.config.get("LOG_FILE")
|
|
or ""
|
|
).strip()
|
|
if log_file:
|
|
base = Path(log_file).resolve().parent
|
|
else:
|
|
base = Path(self.context.database_path).resolve().parent
|
|
self.service_log = _make_service_logger(base, self.context.logger)
|
|
|
|
|
|
def register_realtime(socket_server: SocketIO, context: EngineContext) -> None:
|
|
"""Register Socket.IO event handlers for the Engine runtime."""
|
|
|
|
adapters = EngineRealtimeAdapters(context)
|
|
logger = context.logger.getChild("realtime.quick_jobs")
|
|
tunnel_service = getattr(context, "reverse_tunnel_service", None)
|
|
if tunnel_service is None:
|
|
tunnel_service = ReverseTunnelService(
|
|
context,
|
|
signer=None,
|
|
db_conn_factory=adapters.db_conn_factory,
|
|
socketio=socket_server,
|
|
)
|
|
tunnel_service.start()
|
|
setattr(context, "reverse_tunnel_service", tunnel_service)
|
|
|
|
@socket_server.on("quick_job_result")
|
|
def _handle_quick_job_result(data: Any) -> None:
|
|
if not isinstance(data, dict):
|
|
logger.debug("quick_job_result payload ignored (non-dict): %r", data)
|
|
return
|
|
|
|
job_id_raw = data.get("job_id")
|
|
try:
|
|
job_id = int(job_id_raw)
|
|
except (TypeError, ValueError):
|
|
logger.debug("quick_job_result missing valid job_id: %r", job_id_raw)
|
|
return
|
|
|
|
status = str(data.get("status") or "").strip() or "Failed"
|
|
stdout = _normalize_text(data.get("stdout"))
|
|
stderr = _normalize_text(data.get("stderr"))
|
|
|
|
conn: Optional[sqlite3.Connection] = None
|
|
cursor = None
|
|
broadcast_payload: Optional[Dict[str, Any]] = None
|
|
|
|
ctx_payload = data.get("context")
|
|
context_info: Optional[Dict[str, Any]] = ctx_payload if isinstance(ctx_payload, dict) else None
|
|
|
|
try:
|
|
conn = adapters.db_conn_factory()
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"UPDATE activity_history SET status=?, stdout=?, stderr=? WHERE id=?",
|
|
(status, stdout, stderr, job_id),
|
|
)
|
|
if cursor.rowcount == 0:
|
|
logger.debug("quick_job_result missing activity_history row for job_id=%s", job_id)
|
|
conn.commit()
|
|
|
|
try:
|
|
cursor.execute(
|
|
"SELECT run_id FROM scheduled_job_run_activity WHERE activity_id=?",
|
|
(job_id,),
|
|
)
|
|
link = cursor.fetchone()
|
|
except sqlite3.Error:
|
|
link = None
|
|
|
|
run_id: Optional[int] = None
|
|
scheduled_ts_ctx: Optional[int] = None
|
|
if link:
|
|
try:
|
|
run_id = int(link[0])
|
|
except Exception:
|
|
run_id = None
|
|
|
|
if run_id is None and context_info:
|
|
ctx_run = context_info.get("scheduled_job_run_id") or context_info.get("run_id")
|
|
try:
|
|
if ctx_run is not None:
|
|
run_id = int(ctx_run)
|
|
except (TypeError, ValueError):
|
|
run_id = None
|
|
try:
|
|
if context_info.get("scheduled_ts") is not None:
|
|
scheduled_ts_ctx = int(context_info.get("scheduled_ts"))
|
|
except (TypeError, ValueError):
|
|
scheduled_ts_ctx = None
|
|
|
|
if run_id is not None:
|
|
ts_now = _now_ts()
|
|
try:
|
|
if status.lower() == "running":
|
|
cursor.execute(
|
|
"UPDATE scheduled_job_runs SET status='Running', updated_at=? WHERE id=?",
|
|
(ts_now, run_id),
|
|
)
|
|
else:
|
|
cursor.execute(
|
|
"""
|
|
UPDATE scheduled_job_runs
|
|
SET status=?,
|
|
finished_ts=COALESCE(finished_ts, ?),
|
|
updated_at=?
|
|
WHERE id=?
|
|
""",
|
|
(status, ts_now, ts_now, run_id),
|
|
)
|
|
if scheduled_ts_ctx is not None:
|
|
cursor.execute(
|
|
"UPDATE scheduled_job_runs SET scheduled_ts=COALESCE(scheduled_ts, ?) WHERE id=?",
|
|
(scheduled_ts_ctx, run_id),
|
|
)
|
|
conn.commit()
|
|
adapters.service_log(
|
|
"scheduled_jobs",
|
|
f"scheduled run update run_id={run_id} activity_id={job_id} status={status}",
|
|
)
|
|
except Exception as exc: # pragma: no cover - defensive guard
|
|
logger.debug(
|
|
"quick_job_result failed to update scheduled_job_runs for job_id=%s run_id=%s: %s",
|
|
job_id,
|
|
run_id,
|
|
exc,
|
|
)
|
|
elif context_info:
|
|
adapters.service_log(
|
|
"scheduled_jobs",
|
|
f"scheduled run update skipped (no run_id) activity_id={job_id} status={status} context={context_info}",
|
|
level="WARNING",
|
|
)
|
|
|
|
try:
|
|
cursor.execute(
|
|
"SELECT id, hostname, status FROM activity_history WHERE id=?",
|
|
(job_id,),
|
|
)
|
|
row = cursor.fetchone()
|
|
except sqlite3.Error:
|
|
row = None
|
|
|
|
if row:
|
|
hostname = (row[1] or "").strip()
|
|
if hostname:
|
|
broadcast_payload = {
|
|
"activity_id": int(row[0]),
|
|
"hostname": hostname,
|
|
"status": row[2] or status,
|
|
"change": "updated",
|
|
"source": "quick_job",
|
|
}
|
|
|
|
adapters.service_log(
|
|
"assemblies",
|
|
f"quick_job_result processed job_id={job_id} status={status}",
|
|
)
|
|
except Exception as exc: # pragma: no cover - defensive guard
|
|
logger.warning(
|
|
"quick_job_result handler error for job_id=%s: %s",
|
|
job_id,
|
|
exc,
|
|
exc_info=True,
|
|
)
|
|
finally:
|
|
if cursor is not None:
|
|
try:
|
|
cursor.close()
|
|
except Exception:
|
|
pass
|
|
if conn is not None:
|
|
try:
|
|
conn.close()
|
|
except Exception:
|
|
pass
|
|
|
|
if broadcast_payload:
|
|
try:
|
|
socket_server.emit("device_activity_changed", broadcast_payload)
|
|
except Exception as exc: # pragma: no cover - defensive guard
|
|
logger.debug(
|
|
"Failed to emit device_activity_changed for job_id=%s: %s",
|
|
job_id,
|
|
exc,
|
|
)
|
|
|
|
@socket_server.on("tunnel_bridge_attach")
|
|
def _tunnel_bridge_attach(data: Any) -> Any:
|
|
"""Placeholder operator bridge attach handler (no data channel yet)."""
|
|
|
|
if not isinstance(data, dict):
|
|
return {"error": "invalid_payload"}
|
|
|
|
tunnel_id = str(data.get("tunnel_id") or "").strip()
|
|
operator_id = str(data.get("operator_id") or "").strip() or None
|
|
if not tunnel_id:
|
|
return {"error": "tunnel_id_required"}
|
|
|
|
try:
|
|
tunnel_service.operator_attach(tunnel_id, operator_id)
|
|
except ValueError as exc:
|
|
return {"error": str(exc)}
|
|
except Exception as exc: # pragma: no cover - defensive guard
|
|
logger.debug("tunnel_bridge_attach failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
|
|
return {"error": "bridge_attach_failed"}
|
|
|
|
return {"status": "ok", "tunnel_id": tunnel_id, "operator_id": operator_id or "-"}
|
|
|
|
def _encode_frame(frame: TunnelFrame) -> str:
|
|
return base64.b64encode(frame.encode()).decode("ascii")
|
|
|
|
def _decode_frame_payload(raw: Any) -> TunnelFrame:
|
|
if isinstance(raw, str):
|
|
try:
|
|
raw_bytes = base64.b64decode(raw)
|
|
except Exception:
|
|
raise ValueError("invalid_frame")
|
|
elif isinstance(raw, (bytes, bytearray)):
|
|
raw_bytes = bytes(raw)
|
|
else:
|
|
raise ValueError("invalid_frame")
|
|
return decode_frame(raw_bytes)
|
|
|
|
@socket_server.on("tunnel_operator_send")
|
|
def _tunnel_operator_send(data: Any) -> Any:
|
|
"""Operator -> agent frame enqueue (placeholder queue)."""
|
|
|
|
if not isinstance(data, dict):
|
|
return {"error": "invalid_payload"}
|
|
tunnel_id = str(data.get("tunnel_id") or "").strip()
|
|
frame_raw = data.get("frame")
|
|
if not tunnel_id or frame_raw is None:
|
|
return {"error": "tunnel_id_and_frame_required"}
|
|
try:
|
|
frame = _decode_frame_payload(frame_raw)
|
|
except Exception as exc:
|
|
return {"error": str(exc)}
|
|
|
|
bridge: Optional[TunnelBridge] = tunnel_service.get_bridge(tunnel_id)
|
|
if bridge is None:
|
|
return {"error": "unknown_tunnel"}
|
|
bridge.operator_to_agent(frame)
|
|
return {"status": "ok"}
|
|
|
|
@socket_server.on("tunnel_operator_poll")
|
|
def _tunnel_operator_poll(data: Any) -> Any:
|
|
"""Operator polls queued frames from agent."""
|
|
|
|
tunnel_id = ""
|
|
if isinstance(data, dict):
|
|
tunnel_id = str(data.get("tunnel_id") or "").strip()
|
|
if not tunnel_id:
|
|
return {"error": "tunnel_id_required"}
|
|
bridge: Optional[TunnelBridge] = tunnel_service.get_bridge(tunnel_id)
|
|
if bridge is None:
|
|
return {"error": "unknown_tunnel"}
|
|
|
|
frames = []
|
|
while True:
|
|
frame = bridge.next_for_operator()
|
|
if frame is None:
|
|
break
|
|
frames.append(_encode_frame(frame))
|
|
return {"frames": frames}
|
|
|
|
# WebUI operator bridge namespace for browser clients
|
|
tunnel_namespace = "/tunnel"
|
|
_operator_sessions: Dict[str, str] = {}
|
|
|
|
def _current_operator() -> Optional[str]:
|
|
username = session.get("username")
|
|
if username:
|
|
return str(username)
|
|
auth_header = (request.headers.get("Authorization") or "").strip()
|
|
token = None
|
|
if auth_header.lower().startswith("bearer "):
|
|
token = auth_header.split(" ", 1)[1].strip()
|
|
if not token:
|
|
token = request.cookies.get("borealis_auth")
|
|
return token or None
|
|
|
|
@socket_server.on("join", namespace=tunnel_namespace)
|
|
def _ws_tunnel_join(data: Any) -> Any:
|
|
if not isinstance(data, dict):
|
|
return {"error": "invalid_payload"}
|
|
operator_id = _current_operator()
|
|
if not operator_id:
|
|
return {"error": "unauthorized"}
|
|
tunnel_id = str(data.get("tunnel_id") or "").strip()
|
|
if not tunnel_id:
|
|
return {"error": "tunnel_id_required"}
|
|
bridge = tunnel_service.get_bridge(tunnel_id)
|
|
if bridge is None:
|
|
return {"error": "unknown_tunnel"}
|
|
try:
|
|
tunnel_service.operator_attach(tunnel_id, operator_id)
|
|
except Exception as exc:
|
|
logger.debug("ws_tunnel_join failed tunnel_id=%s: %s", tunnel_id, exc, exc_info=True)
|
|
return {"error": "attach_failed"}
|
|
sid = request.sid
|
|
_operator_sessions[sid] = tunnel_id
|
|
return {"status": "ok", "tunnel_id": tunnel_id}
|
|
|
|
@socket_server.on("send", namespace=tunnel_namespace)
|
|
def _ws_tunnel_send(data: Any) -> Any:
|
|
sid = request.sid
|
|
tunnel_id = _operator_sessions.get(sid)
|
|
if not tunnel_id:
|
|
return {"error": "not_joined"}
|
|
if not isinstance(data, dict):
|
|
return {"error": "invalid_payload"}
|
|
frame_raw = data.get("frame")
|
|
if frame_raw is None:
|
|
return {"error": "frame_required"}
|
|
try:
|
|
frame = _decode_frame_payload(frame_raw)
|
|
except Exception:
|
|
return {"error": "invalid_frame"}
|
|
bridge = tunnel_service.get_bridge(tunnel_id)
|
|
if bridge is None:
|
|
return {"error": "unknown_tunnel"}
|
|
bridge.operator_to_agent(frame)
|
|
return {"status": "ok"}
|
|
|
|
@socket_server.on("poll", namespace=tunnel_namespace)
|
|
def _ws_tunnel_poll() -> Any:
|
|
sid = request.sid
|
|
tunnel_id = _operator_sessions.get(sid)
|
|
if not tunnel_id:
|
|
return {"error": "not_joined"}
|
|
bridge = tunnel_service.get_bridge(tunnel_id)
|
|
if bridge is None:
|
|
return {"error": "unknown_tunnel"}
|
|
frames = []
|
|
while True:
|
|
frame = bridge.next_for_operator()
|
|
if frame is None:
|
|
break
|
|
frames.append(_encode_frame(frame))
|
|
return {"frames": frames}
|
|
|
|
@socket_server.on("disconnect", namespace=tunnel_namespace)
|
|
def _ws_tunnel_disconnect():
|
|
sid = request.sid
|
|
_operator_sessions.pop(sid, None)
|