Refactor WebSocket console code to work with FastAPI.

Fix endpoint routes.

gns3server/endpoints/compute/atm_switch_nodes.py

@@ -47,7 +47,7 @@ async def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.ATMSwitch,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create ATM switch node"}})

gns3server/endpoints/compute/cloud_nodes.py

@@ -48,7 +48,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.Cloud,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create cloud node"}})

gns3server/endpoints/compute/docker_nodes.py

@@ -21,7 +21,7 @@ API endpoints for Docker nodes.
 
 import os
 
-from fastapi import APIRouter, Depends, Body, status
+from fastapi import APIRouter, WebSocket, Depends, Body, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse
 from uuid import UUID
@@ -47,7 +47,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.Docker,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create Docker node"}})
@@ -290,14 +290,6 @@ async def stop_capture(adapter_number: int, port_number: int, node: DockerVM = D
     await node.stop_capture(adapter_number)
 
 
-@router.post("/{node_id}/console/reset",
-             status_code=status.HTTP_204_NO_CONTENT,
-             responses=responses)
-async def reset_console(node: DockerVM = Depends(dep_node)):
-
-    await node.reset_console()
-
-
 @router.get("/{node_id}/adapters/{adapter_number}/ports/{port_number}/pcap",
             responses=responses)
 async def stream_pcap_file(adapter_number: int, port_number: int, node: DockerVM = Depends(dep_node)):
@@ -310,18 +302,19 @@ async def stream_pcap_file(adapter_number: int, port_number: int, node: DockerVM
     stream = Docker.instance().stream_pcap_file(nio, node.project.id)
     return StreamingResponse(stream, media_type="application/vnd.tcpdump.pcap")
 
-# @Route.get(
-#     r"/projects/{project_id}/docker/nodes/{node_id}/console/ws",
-#     description="WebSocket for console",
-#     parameters={
-#         "project_id": "Project UUID",
-#         "node_id": "Node UUID",
-#     })
-# async def console_ws(request, response):
-#
-#     docker_manager = Docker.instance()
-#     container = docker_manager.get_node(request.match_info["node_id"], project_id=request.match_info["project_id"])
-#     return await container.start_websocket_console(request)
+
+@router.websocket("/{node_id}/console/ws")
+async def console_ws(websocket: WebSocket, node: DockerVM = Depends(dep_node)):
+    """
+    Console WebSocket.
+    """
+
+    await node.start_websocket_console(websocket)
 
 
+@router.post("/{node_id}/console/reset",
+             status_code=status.HTTP_204_NO_CONTENT,
+             responses=responses)
+async def reset_console(node: DockerVM = Depends(dep_node)):
 
+    await node.reset_console()

gns3server/endpoints/compute/dynamips_nodes.py

@@ -22,7 +22,7 @@ API endpoints for Dynamips nodes.
 import os
 import sys
 
-from fastapi import APIRouter, Depends, status
+from fastapi import APIRouter, WebSocket, Depends, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse
 from typing import List
@@ -56,7 +56,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.Dynamips,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create Dynamips node"}})
@@ -299,18 +299,13 @@ async def duplicate_router(destination_node_id: UUID, node: Router = Depends(dep
     return new_node.__json__()
 
 
-# @Route.get(
-#     r"/projects/{project_id}/dynamips/nodes/{node_id}/console/ws",
-#     description="WebSocket for console",
-#     parameters={
-#         "project_id": "Project UUID",
-#         "node_id": "Node UUID",
-#     })
-# async def console_ws(request, response):
-#
-#     dynamips_manager = Dynamips.instance()
-#     vm = dynamips_manager.get_node(request.match_info["node_id"], project_id=request.match_info["project_id"])
-#     return await vm.start_websocket_console(request)
+@router.websocket("/{node_id}/console/ws")
+async def console_ws(websocket: WebSocket, node: Router = Depends(dep_node)):
+    """
+    Console WebSocket.
+    """
+
+    await node.start_websocket_console(websocket)
 
 
 @router.post("/{node_id}/console/reset",

gns3server/endpoints/compute/ethernet_hub_nodes.py

@@ -47,7 +47,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.EthernetHub,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create Ethernet hub node"}})

gns3server/endpoints/compute/ethernet_switch_nodes.py

@@ -47,7 +47,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.EthernetSwitch,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create Ethernet switch node"}})

gns3server/endpoints/compute/frame_relay_switch_nodes.py

@@ -47,7 +47,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.FrameRelaySwitch,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create Frame Relay switch node"}})

gns3server/endpoints/compute/iou_nodes.py

@@ -21,7 +21,7 @@ API endpoints for IOU nodes.
 
 import os
 
-from fastapi import APIRouter, Depends, Body, status
+from fastapi import APIRouter, WebSocket, Depends, Body, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse
 from typing import Union
@@ -48,7 +48,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.IOU,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create IOU node"}})
@@ -275,23 +275,18 @@ async def stream_pcap_file(adapter_number: int, port_number: int, node: IOUVM =
     return StreamingResponse(stream, media_type="application/vnd.tcpdump.pcap")
 
 
+@router.websocket("/{node_id}/console/ws")
+async def console_ws(websocket: WebSocket, node: IOUVM = Depends(dep_node)):
+    """
+    Console WebSocket.
+    """
+
+    await node.start_websocket_console(websocket)
+
+
 @router.post("/{node_id}/console/reset",
              status_code=status.HTTP_204_NO_CONTENT,
              responses=responses)
 async def reset_console(node: IOUVM = Depends(dep_node)):
 
     await node.reset_console()
-
-
-# @Route.get(
-#     r"/projects/{project_id}/iou/nodes/{node_id}/console/ws",
-#     description="WebSocket for console",
-#     parameters={
-#         "project_id": "Project UUID",
-#         "node_id": "Node UUID",
-#     })
-# async def console_ws(request, response):
-#
-#     iou_manager = IOU.instance()
-#     vm = iou_manager.get_node(request.match_info["node_id"], project_id=request.match_info["project_id"])
-#     return await vm.start_websocket_console(request)

gns3server/endpoints/compute/nat_nodes.py

@@ -48,7 +48,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.NAT,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create NAT node"}})

gns3server/endpoints/compute/notifications.py

@@ -19,11 +19,10 @@
 API endpoints for compute notifications.
 """
 
-from fastapi import APIRouter, WebSocket, WebSocketDisconnect
-from websockets.exceptions import WebSocketException
-from typing import List
-
+import asyncio
+from fastapi import APIRouter, WebSocket
 from gns3server.compute.notification_manager import NotificationManager
+from starlette.endpoints import WebSocketEndpoint
 
 import logging
 log = logging.getLogger(__name__)
@@ -31,48 +30,63 @@ log = logging.getLogger(__name__)
 router = APIRouter()
 
 
-class ConnectionManager:
-    def __init__(self):
-        self.active_connections: List[WebSocket] = []
+@router.websocket_route("/notifications/ws")
+class ComputeWebSocketNotifications(WebSocketEndpoint):
+    """
+    Receive compute notifications about the controller from WebSocket stream.
+    """
+
+    async def on_connect(self, websocket: WebSocket) -> None:
 
-    async def connect(self, websocket: WebSocket):
         await websocket.accept()
-        self.active_connections.append(websocket)
+        log.info(f"New client {websocket.client.host}:{websocket.client.port} has connected to compute WebSocket")
+        self._notification_task = asyncio.ensure_future(self._stream_notifications(websocket))
 
-    def disconnect(self, websocket: WebSocket):
+    async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:
 
-        self.active_connections.remove(websocket)
+        self._notification_task.cancel()
+        log.info(f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller WebSocket"
+                 f" with close code {close_code}")
 
-    async def close_active_connections(self):
+    async def _stream_notifications(self, websocket: WebSocket) -> None:
 
-        for websocket in self.active_connections:
-            await websocket.close()
-
-    async def send_text(self, message: str, websocket: WebSocket):
-        await websocket.send_text(message)
-
-    async def broadcast(self, message: str):
-        for connection in self.active_connections:
-            await connection.send_text(message)
-
-
-manager = ConnectionManager()
-
-
-@router.websocket("/notifications/ws")
-async def compute_notifications(websocket: WebSocket):
-
-    log.info("Client has disconnected from compute WebSocket")
-    notifications = NotificationManager.instance()
-    await manager.connect(websocket)
-    try:
-        log.info("New client has connected to compute WebSocket")
-        with notifications.queue() as queue:
+        with NotificationManager.instance().queue() as queue:
             while True:
                 notification = await queue.get_json(5)
-                await manager.send_text(notification, websocket)
-    except (WebSocketException, WebSocketDisconnect) as e:
-        log.info("Client has disconnected from compute WebSocket: {}".format(e))
-    finally:
-        await websocket.close()
-        manager.disconnect(websocket)
+                await websocket.send_text(notification)
+
+
+if __name__ == '__main__':
+
+    import uvicorn
+    from fastapi import FastAPI
+    from starlette.responses import HTMLResponse
+
+    app = FastAPI()
+    app.include_router(router)
+
+    html = """
+    <!DOCTYPE html>
+    <html>
+        <body>
+            <ul id='messages'>
+            </ul>
+            <script>
+                var ws = new WebSocket("ws://localhost:8000/notifications/ws");        
+                ws.onmessage = function(event) {
+                    var messages = document.getElementById('messages')
+                    var message = document.createElement('li')
+                    var content = document.createTextNode(event.data)
+                    message.appendChild(content)
+                    messages.appendChild(message)
+                };
+            </script>
+        </body>
+    </html>
+    """
+
+    @app.get("/")
+    async def get() -> HTMLResponse:
+        return HTMLResponse(html)
+
+    uvicorn.run(app, host="localhost", port=8000)

gns3server/endpoints/compute/projects.py

@@ -19,8 +19,6 @@
 API endpoints for projects.
 """
 
-import shutil
-import aiohttp
 import os
 
 import logging
@@ -108,6 +106,7 @@ async def close_project(project: Project = Depends(dep_project)):
     Close a project on the compute.
     """
 
+    # FIXME
     if _notifications_listening.setdefault(project.id, 0) <= 1:
         await project.close()
         ProjectManager.instance().remove_project(project.id)
@@ -234,6 +233,6 @@ async def write_file(file_path: str, request: Request, project: Project = Depend
             pass  # FIXME
 
     except FileNotFoundError:
-        raise aiohttp.web.HTTPNotFound()
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
     except PermissionError:
-        raise aiohttp.web.HTTPForbidden()
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)

gns3server/endpoints/compute/qemu_nodes.py

@@ -22,14 +22,13 @@ API endpoints for Qemu nodes.
 import os
 import sys
 
-from fastapi import APIRouter, Depends, Body, status
+from fastapi import APIRouter, WebSocket, Depends, Body, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse
 from uuid import UUID
 
 from gns3server.endpoints import schemas
 from gns3server.compute.project_manager import ProjectManager
-from gns3server.compute.compute_error import ComputeError
 from gns3server.compute.qemu import Qemu
 from gns3server.compute.qemu.qemu_vm import QemuVM
 
@@ -50,7 +49,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.Qemu,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create Qemu node"}})
@@ -281,14 +280,6 @@ async def stop_capture(adapter_number: int, port_number: int, node: QemuVM = Dep
     await node.stop_capture(adapter_number)
 
 
-@router.post("/{node_id}/console/reset",
-             status_code=status.HTTP_204_NO_CONTENT,
-             responses=responses)
-async def reset_console(node: QemuVM = Depends(dep_node)):
-
-    await node.reset_console()
-
-
 @router.get("/{node_id}/adapters/{adapter_number}/ports/{port_number}/pcap",
             responses=responses)
 async def stream_pcap_file(adapter_number: int, port_number: int, node: QemuVM = Depends(dep_node)):
@@ -302,16 +293,18 @@ async def stream_pcap_file(adapter_number: int, port_number: int, node: QemuVM =
     return StreamingResponse(stream, media_type="application/vnd.tcpdump.pcap")
 
 
-# @Route.get(
-#     r"/projects/{project_id}/qemu/nodes/{node_id}/console/ws",
-#     description="WebSocket for console",
-#     parameters={
-#         "project_id": "Project UUID",
-#         "node_id": "Node UUID",
-#     })
-# async def console_ws(request, response):
-#
-#     qemu_manager = Qemu.instance()
-#     vm = qemu_manager.get_node(request.match_info["node_id"], project_id=request.match_info["project_id"])
-#     return await vm.start_websocket_console(request)
+@router.websocket("/{node_id}/console/ws")
+async def console_ws(websocket: WebSocket, node: QemuVM = Depends(dep_node)):
+    """
+    Console WebSocket.
+    """
 
+    await node.start_websocket_console(websocket)
+
+
+@router.post("/{node_id}/console/reset",
+             status_code=status.HTTP_204_NO_CONTENT,
+             responses=responses)
+async def reset_console(node: QemuVM = Depends(dep_node)):
+
+    await node.reset_console()

gns3server/endpoints/compute/virtualbox_nodes.py

@@ -21,7 +21,7 @@ API endpoints for VirtualBox nodes.
 
 import os
 
-from fastapi import APIRouter, Depends, status
+from fastapi import APIRouter, WebSocket, Depends, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse
 from uuid import UUID
@@ -49,7 +49,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.VirtualBox,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create VirtualBox node"}})
@@ -288,14 +288,6 @@ async def stop_capture(adapter_number: int, port_number: int, node: VirtualBoxVM
     await node.stop_capture(adapter_number)
 
 
-@router.post("/{node_id}/console/reset",
-             status_code=status.HTTP_204_NO_CONTENT,
-             responses=responses)
-async def reset_console(node: VirtualBoxVM = Depends(dep_node)):
-
-    await node.reset_console()
-
-
 @router.get("/{node_id}/adapters/{adapter_number}/ports/{port_number}/pcap",
             responses=responses)
 async def stream_pcap_file(adapter_number: int, port_number: int, node: VirtualBoxVM = Depends(dep_node)):
@@ -309,15 +301,18 @@ async def stream_pcap_file(adapter_number: int, port_number: int, node: VirtualB
     return StreamingResponse(stream, media_type="application/vnd.tcpdump.pcap")
 
 
-# @Route.get(
-#     r"/projects/{project_id}/virtualbox/nodes/{node_id}/console/ws",
-#     description="WebSocket for console",
-#     parameters={
-#         "project_id": "Project UUID",
-#         "node_id": "Node UUID",
-#     })
-# async def console_ws(request, response):
-#
-#     virtualbox_manager = VirtualBox.instance()
-#     vm = virtualbox_manager.get_node(request.match_info["node_id"], project_id=request.match_info["project_id"])
-#     return await vm.start_websocket_console(request)
+@router.websocket("/{node_id}/console/ws")
+async def console_ws(websocket: WebSocket, node: VirtualBoxVM = Depends(dep_node)):
+    """
+    Console WebSocket.
+    """
+
+    await node.start_websocket_console(websocket)
+
+
+@router.post("/{node_id}/console/reset",
+             status_code=status.HTTP_204_NO_CONTENT,
+             responses=responses)
+async def reset_console(node: VirtualBoxVM = Depends(dep_node)):
+
+    await node.reset_console()

gns3server/endpoints/compute/vmware_nodes.py

@@ -21,7 +21,7 @@ API endpoints for VMware nodes.
 
 import os
 
-from fastapi import APIRouter, Depends, status
+from fastapi import APIRouter, WebSocket, Depends, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse
 from uuid import UUID
@@ -48,7 +48,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.VMware,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create VMware node"}})
@@ -253,14 +253,6 @@ async def stop_capture(adapter_number: int, port_number: int, node: VMwareVM = D
     await node.stop_capture(adapter_number)
 
 
-@router.post("/{node_id}/console/reset",
-             status_code=status.HTTP_204_NO_CONTENT,
-             responses=responses)
-async def reset_console(node: VMwareVM = Depends(dep_node)):
-
-    await node.reset_console()
-
-
 @router.get("/{node_id}/adapters/{adapter_number}/ports/{port_number}/pcap",
             responses=responses)
 async def stream_pcap_file(adapter_number: int, port_number: int, node: VMwareVM = Depends(dep_node)):
@@ -289,16 +281,18 @@ def allocate_vmnet(node: VMwareVM = Depends(dep_node)) -> dict:
     return {"vmnet": vmnet}
 
 
-# @Route.get(
-#     r"/projects/{project_id}/vmware/nodes/{node_id}/console/ws",
-#     description="WebSocket for console",
-#     parameters={
-#         "project_id": "Project UUID",
-#         "node_id": "Node UUID",
-#     })
-# async def console_ws(request, response):
-#
-#     vmware_manager = VMware.instance()
-#     vm = vmware_manager.get_node(request.match_info["node_id"], project_id=request.match_info["project_id"])
-#     return await vm.start_websocket_console(request)
-#
+@router.websocket("/{node_id}/console/ws")
+async def console_ws(websocket: WebSocket, node: VMwareVM = Depends(dep_node)):
+    """
+    Console WebSocket.
+    """
+
+    await node.start_websocket_console(websocket)
+
+
+@router.post("/{node_id}/console/reset",
+             status_code=status.HTTP_204_NO_CONTENT,
+             responses=responses)
+async def reset_console(node: VMwareVM = Depends(dep_node)):
+
+    await node.reset_console()

gns3server/endpoints/compute/vpcs_nodes.py

@@ -21,14 +21,13 @@ API endpoints for VPCS nodes.
 
 import os
 
-from fastapi import APIRouter, Depends, Body, status
+from fastapi import APIRouter, WebSocket, Depends, Body, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse
 from uuid import UUID
 
 from gns3server.endpoints import schemas
 from gns3server.compute.vpcs import VPCS
-from gns3server.compute.project_manager import ProjectManager
 from gns3server.compute.vpcs.vpcs_vm import VPCSVM
 
 router = APIRouter()
@@ -48,7 +47,7 @@ def dep_node(project_id: UUID, node_id: UUID):
     return node
 
 
-@router.post("/",
+@router.post("",
              response_model=schemas.VPCS,
              status_code=status.HTTP_201_CREATED,
              responses={409: {"model": schemas.ErrorMessage, "description": "Could not create VMware node"}})
@@ -258,15 +257,10 @@ async def stream_pcap_file(adapter_number: int, port_number: int, node: VPCSVM =
     return StreamingResponse(stream, media_type="application/vnd.tcpdump.pcap")
 
 
-# @Route.get(
-#     r"/projects/{project_id}/vpcs/nodes/{node_id}/console/ws",
-#     description="WebSocket for console",
-#     parameters={
-#         "project_id": "Project UUID",
-#         "node_id": "Node UUID",
-#     })
-# async def console_ws(request, response):
-#
-#     vpcs_manager = VPCS.instance()
-#     vm = vpcs_manager.get_node(request.match_info["node_id"], project_id=request.match_info["project_id"])
-#     return await vm.start_websocket_console(request)
+@router.websocket("/{node_id}/console/ws")
+async def console_ws(websocket: WebSocket, node: VPCSVM = Depends(dep_node)):
+    """
+    Console WebSocket.
+    """
+
+    await node.start_websocket_console(websocket)