fix(backend): add vnc-port and screenshot endpoints to slave agent

Author: lg-epitech

backend/app/routes/slave_agent.py

@@ -2,10 +2,12 @@
 
 All endpoints require X-Slave-Token authentication.
 """
-from fastapi import APIRouter, Depends, status
+from fastapi import APIRouter, Depends, Response, status
 from app.models.vm import VmCreate, VmRead
 from app.services.vm_service import VmService
 from app.services.host_service import HostService
+from app.services.vm_screenshot import capture_screenshot
+from app.utils.vnc import get_vnc_port
 from app.utils.slave_auth import require_slave_token
 
 router = APIRouter()
@@ -60,6 +62,30 @@ def delete_vm(vm_id: str):
     VmService.remove_vm(vm_id)
 
 
+@router.get(
+    "/vms/{vm_id}/vnc-port",
+    status_code=status.HTTP_200_OK,
+    dependencies=[Depends(require_slave_token)],
+)
+def get_vm_vnc_port(vm_id: str):
+    port = get_vnc_port(vm_id)
+    return {"port": port}
+
+
+@router.get(
+    "/vms/{vm_id}/screenshot",
+    status_code=status.HTTP_200_OK,
+    dependencies=[Depends(require_slave_token)],
+)
+def get_vm_screenshot(vm_id: str):
+    jpeg_bytes = capture_screenshot(vm_id)
+    return Response(
+        content=jpeg_bytes,
+        media_type="image/jpeg",
+        headers={"Cache-Control": "no-store"},
+    )
+
+
 @router.get(
     "/host/info",
     status_code=status.HTTP_200_OK,

backend/app/routes/tunnel.py

@@ -2,6 +2,7 @@
 from app.utils.crypto import decrypt_secret
 from app.utils.auth import decode_access_token, user_has_policy
 from app.services.guacamole import build_instruction, guacd_handshake, read_instruction
+from app.services.vm_service import VmService
 from app.orm.vm_credential import VmCredentialORM
 from app.orm.user import UserORM
 from app.core.config import engine, GUACD_HOST, GUACD_PORT, VNC_HOST
@@ -144,47 +145,65 @@ async def vm_tunnel(
         await websocket.close(code=4001, reason="Provide credential or vm_id+token")
         return
 
-    try:
-        vnc_port = await asyncio.to_thread(get_vnc_port, resolved_vm_id)
-    except Exception as exc:
-        detail = getattr(exc, "detail", str(exc))
-        await websocket.close(code=4002, reason=detail)
-        return
+    slave = await asyncio.to_thread(VmService._get_slave_for_vm, resolved_vm_id)
+
+    if slave:
+        from app.services.slave_client import slave_get_vnc_port
+        try:
+            vnc_port = await asyncio.to_thread(slave_get_vnc_port, slave, resolved_vm_id)
+        except Exception as exc:
+            detail = getattr(exc, "detail", str(exc))
+            await websocket.close(code=4002, reason=detail)
+            return
+        guacd_host = slave.hostname
+        guacd_port = GUACD_PORT
+        vnc_host = "127.0.0.1"
+    else:
+        try:
+            vnc_port = await asyncio.to_thread(get_vnc_port, resolved_vm_id)
+        except Exception as exc:
+            detail = getattr(exc, "detail", str(exc))
+            await websocket.close(code=4002, reason=detail)
+            return
+        guacd_host = GUACD_HOST
+        guacd_port = GUACD_PORT
+        vnc_host = VNC_HOST
 
     logger.warning(
-        "Tunnel: vm_id=%s vnc=%s:%s guacd=%s:%s",
-        resolved_vm_id, VNC_HOST, vnc_port, GUACD_HOST, GUACD_PORT,
+        "Tunnel: vm_id=%s vnc=%s:%s guacd=%s:%s slave=%s",
+        resolved_vm_id, vnc_host, vnc_port, guacd_host, guacd_port,
+        slave.hostname if slave else None,
     )
 
     await websocket.accept(subprotocol="guacamole")
 
     try:
-        reader, writer = await asyncio.open_connection(GUACD_HOST, GUACD_PORT)
+        reader, writer = await asyncio.open_connection(guacd_host, guacd_port)
     except Exception as exc:
         await websocket.close(
             code=1011,
-            reason=f"Cannot connect to guacd {GUACD_HOST}:{GUACD_PORT}: {exc}",
+            reason=f"Cannot connect to guacd {guacd_host}:{guacd_port}: {exc}",
         )
         return
 
     try:
         first_instruction = await guacd_handshake(
             reader,
             writer,
-            vnc_host=VNC_HOST,
+            vnc_host=vnc_host,
             vnc_port=vnc_port,
             width=width,
             height=height,
         )
         logger.warning(
-            "Tunnel connected via configured vnc host=%s port=%s",
-            VNC_HOST,
+            "Tunnel connected via vnc host=%s port=%s",
+            vnc_host,
             vnc_port,
         )
     except Exception as exc:
         logger.warning(
-            "Tunnel VNC connect failed for configured host=%s port=%s: %s",
-            VNC_HOST,
+            "Tunnel VNC connect failed for host=%s port=%s: %s",
+            vnc_host,
             vnc_port,
             exc,
         )

backend/app/routes/vm.py

@@ -38,7 +38,12 @@ async def get_vm_screenshot(vm_id: str, token: str = Query(...)):
         if not user_has_policy(user, "vms:get"):
             raise HTTPException(status.HTTP_403_FORBIDDEN,
                                 "Missing vms:get policy")
-    jpeg_bytes = await asyncio.to_thread(capture_screenshot, vm_id)
+    slave = await asyncio.to_thread(VmService._get_slave_for_vm, vm_id)
+    if slave:
+        from app.services.slave_client import slave_get_screenshot
+        jpeg_bytes = await asyncio.to_thread(slave_get_screenshot, slave, vm_id)
+    else:
+        jpeg_bytes = await asyncio.to_thread(capture_screenshot, vm_id)
     return Response(
         content=jpeg_bytes,
         media_type="image/jpeg",

backend/app/services/slave_client.py

@@ -61,6 +61,22 @@ def slave_delete_vm(slave: SlaveORM, vm_id: str) -> dict:
     return slave_request(slave, "DELETE", f"/vms/{vm_id}")
 
 
+def slave_get_vnc_port(slave: SlaveORM, vm_id: str) -> int:
+    """Get the VNC port for a VM on a slave node."""
+    data = slave_request(slave, "GET", f"/vms/{vm_id}/vnc-port")
+    return data["port"]
+
+
+def slave_get_screenshot(slave: SlaveORM, vm_id: str) -> bytes:
+    """Get a screenshot from a VM on a slave node."""
+    url = f"{_slave_base_url(slave)}/vms/{vm_id}/screenshot"
+    headers = _slave_headers(slave)
+    with httpx.Client(timeout=TIMEOUT) as client:
+        response = client.get(url, headers=headers)
+        response.raise_for_status()
+        return response.content
+
+
 def slave_get_host_info(slave: SlaveORM) -> dict:
     """Get host info from a slave node."""
     return slave_request(slave, "GET", "/host/info")