from __future__ import annotations import json import logging from datetime import timedelta from fastapi import APIRouter, WebSocket, WebSocketDisconnect from sqlalchemy import select from starlette.concurrency import run_in_threadpool from ..api.computers import METRIC_HISTORY, serialize as ser_computer from ..api.deps import computer_from_device_token, principal_from_token from ..config import settings from ..database import SessionLocal from ..models import Computer, ComputerSession, Student, utcnow from ..security import build_signed_command from ..services import audit, commands as cmdsvc from ..services.notify import notify from ..services.policy import effective_policy from ..services.settings_service import get_int from .manager import DashClient, hub log = logging.getLogger("cms.ws") router = APIRouter() def _comp_payload(db, comp): st = db.execute(select(Student).where(Student.computer_id == comp.id)).scalars().first() return ser_computer(comp, st) def _on_connect(token: str): with SessionLocal() as db: comp = computer_from_device_token(db, token) if not comp: return None comp.is_online, comp.last_seen_at, comp.usage_state = True, utcnow(), comp.usage_state if comp.usage_state != "offline" else "idle" db.commit() pol = effective_policy(db, comp) hub.emit({"type": "computer_update", "data": _comp_payload(db, comp)}, classroom_id=comp.classroom_id) audit.log(db, "agent", comp.hostname, "agent.online", computer=comp) return comp.id, comp.device_uid, pol def _on_disconnect(cid: int): with SessionLocal() as db: comp = db.get(Computer, cid) if not comp or not comp.is_online: return comp.is_online, comp.usage_state = False, "offline" for s in db.execute(select(ComputerSession).where(ComputerSession.computer_id == cid, ComputerSession.ended_at.is_(None))).scalars(): s.ended_at = utcnow() db.commit() notify(db, "computer_offline", f"Komputer offline: {comp.label or comp.hostname}", None, severity="warning", computer=comp) hub.emit({"type": "computer_update", "data": _comp_payload(db, comp)}, classroom_id=comp.classroom_id) def _handle_agent(cid: int, msg: dict): t = msg.get("type") with SessionLocal() as db: comp = db.get(Computer, cid) if not comp: return if t == "heartbeat": m = msg.get("metrics") or {} comp.last_seen_at, comp.is_online = utcnow(), True comp.ip_address = msg.get("ip") or comp.ip_address comp.mac_address = msg.get("mac") or comp.mac_address comp.windows_user = msg.get("user") or comp.windows_user comp.agent_version = msg.get("agent_version") or comp.agent_version new_state = msg.get("state") if msg.get("state") in ("idle", "in_use", "locked") else "idle" old_state = comp.usage_state comp.usage_state = new_state comp.metrics = {**m, "processes": (msg.get("processes") or [])[:40], "browsers": (msg.get("browsers") or [])[:10], "policy_version": msg.get("policy_version"), "dry_run": bool(msg.get("dry_run"))} if new_state == "in_use" and old_state != "in_use": db.add(ComputerSession(computer_id=cid, kind="usage", windows_user=comp.windows_user)) elif new_state != "in_use" and old_state == "in_use": for s in db.execute(select(ComputerSession).where(ComputerSession.computer_id == cid, ComputerSession.kind == "usage", ComputerSession.ended_at.is_(None))).scalars(): s.ended_at = utcnow() db.commit() h = METRIC_HISTORY.setdefault(cid, []) h.append({"t": int(utcnow().timestamp()), "cpu": m.get("cpu"), "ram": m.get("ram_percent")}) del h[:-120] hub.emit({"type": "computer_update", "data": _comp_payload(db, comp)}, classroom_id=comp.classroom_id) want = effective_policy(db, comp) if msg.get("policy_version") != want["version"]: hub.send_agent_threadsafe(cid, {"type": "policy_update", "policy": want}) elif t == "command_result": cmdsvc.finish_from_agent(db, int(msg.get("id", 0)), comp, "success" if msg.get("ok") else "failed", str(msg.get("detail", ""))) elif t == "policy_applied": comp.policy_hash, comp.policy_applied_at = msg.get("version"), utcnow() db.commit() if msg.get("errors"): notify(db, "policy_error", f"Kebijakan gagal sebagian di {comp.label or comp.hostname}", "; ".join(map(str, msg["errors"]))[:300], severity="error", computer=comp) audit.log(db, "agent", comp.hostname, "policy.apply", computer=comp, status="failed", detail={"errors": msg["errors"]}) else: audit.log(db, "agent", comp.hostname, "policy.apply", computer=comp, detail={"version": msg.get("version")}) elif t == "inventory_hint": pass @router.websocket("/ws/agent") async def ws_agent(ws: WebSocket): auth = ws.headers.get("authorization", "") token = auth[7:].strip() if auth.lower().startswith("bearer ") else "" info = await run_in_threadpool(_on_connect, token) if token else None if not info: await ws.close(code=4401) return cid, uid, pol = info await ws.accept() old = hub.agents.get(cid) hub.agents[cid] = ws if old is not None: try: await old.close(code=4000) except Exception: pass await ws.send_text(json.dumps({"type": "policy_update", "policy": pol})) try: while True: raw = await ws.receive_text() if len(raw) > 4_000_000: continue msg = json.loads(raw) if msg.get("type") == "screen_frame": await hub.relay_screen(cid, {"jpeg": msg.get("jpeg"), "w": msg.get("w"), "h": msg.get("h")}) elif msg.get("type") == "control_state": for d in hub.viewers(cid): await d.ws.send_text(json.dumps({"type": "control_state", "computer_id": cid, "state": msg.get("state")})) else: await run_in_threadpool(_handle_agent, cid, msg) except (WebSocketDisconnect, json.JSONDecodeError): pass except Exception as e: # pragma: no cover log.warning("agent ws error: %s", e) finally: if hub.agents.get(cid) is ws: hub.agents.pop(cid, None) for d in hub.viewers(cid): d.screen_computer = None await run_in_threadpool(_on_disconnect, cid) # ---- Dashboard --------------------------------------------------------------------------------------- def _dash_auth(token: str): with SessionLocal() as db: return principal_from_token(db, token) def _screen_start(pr, cid: int, viewer: str): with SessionLocal() as db: comp = db.get(Computer, cid) if not comp or not pr.can_see(comp) or not pr.can("screen"): return None fps, q, mw = get_int(db, "screen_fps", 2), get_int(db, "screen_quality", 50), get_int(db, "screen_max_width", 1280) sess = ComputerSession(computer_id=cid, kind="screen_view", started_by=pr.name) db.add(sess) db.commit() audit.log(db, pr.role, pr.name, "screen.start", actor_id=pr.id, computer=comp) signed = build_signed_command(comp.device_uid, 0, "screen_start", {"session_id": sess.id, "viewer": viewer, "fps": fps, "quality": q, "max_width": mw}, settings.command_ttl_seconds) return signed, comp.device_uid def _screen_stop(pr, cid: int): with SessionLocal() as db: comp = db.get(Computer, cid) if not comp: return None for s in db.execute(select(ComputerSession).where(ComputerSession.computer_id == cid, ComputerSession.kind == "screen_view", ComputerSession.ended_at.is_(None))).scalars(): s.ended_at = utcnow() db.commit() audit.log(db, pr.role, pr.name, "screen.stop", actor_id=pr.id, computer=comp) return build_signed_command(comp.device_uid, 0, "screen_stop", {}, settings.command_ttl_seconds) def _signed(pr, cid: int, action: str, params: dict, perm: str = "screen"): with SessionLocal() as db: comp = db.get(Computer, cid) if not comp or not pr.can_see(comp) or not pr.can(perm): return None if action == "control_request": audit.log(db, pr.role, pr.name, "remote_control.request", actor_id=pr.id, computer=comp) return build_signed_command(comp.device_uid, 0, action, params, 60 if action == "control_input" else settings.command_ttl_seconds) async def _stop_view(d: DashClient): cid = d.screen_computer d.screen_computer = None if cid is not None and not hub.viewers(cid): signed = await run_in_threadpool(_screen_stop, _FakePr(d), cid) if signed: await hub.send_agent(cid, {"type": "command", **signed}) class _FakePr: def __init__(self, d: DashClient): self.role, self.id, self.name = d.role, d.uid, d.name @router.websocket("/ws/dashboard") async def ws_dashboard(ws: WebSocket): token = ws.query_params.get("token", "") pr = await run_in_threadpool(_dash_auth, token) if not pr: await ws.close(code=4401) return await ws.accept() d = DashClient(ws=ws, role=pr.role, uid=pr.id, name=pr.name, classroom_ids=pr.classroom_ids) hub.dashboards.append(d) try: while True: msg = json.loads(await ws.receive_text()) t, cid = msg.get("type"), msg.get("computer_id") if t == "ping": await ws.send_text('{"type":"pong"}') elif t == "screen_subscribe" and isinstance(cid, int): await _stop_view(d) res = await run_in_threadpool(_screen_start, pr, cid, pr.name) if not res or not hub.is_online(cid): await ws.send_text(json.dumps({"type": "screen_error", "computer_id": cid, "message": "Tidak diizinkan atau komputer offline"})) continue d.screen_computer = cid await hub.send_agent(cid, {"type": "command", **res[0]}) await ws.send_text(json.dumps({"type": "screen_started", "computer_id": cid})) elif t == "screen_unsubscribe": await _stop_view(d) elif t in ("control_request", "control_stop") and isinstance(cid, int): signed = await run_in_threadpool(_signed, pr, cid, t, {"viewer": pr.name}) if signed: await hub.send_agent(cid, {"type": "command", **signed}) elif t == "control_input" and isinstance(cid, int) and d.screen_computer == cid: ev = msg.get("event") or {} signed = await run_in_threadpool(_signed, pr, cid, "control_input", {"event": ev}) if signed: await hub.send_agent(cid, {"type": "command", **signed}) except (WebSocketDisconnect, json.JSONDecodeError): pass finally: await _stop_view(d) if d in hub.dashboards: hub.dashboards.remove(d)