gns3-server/gns3server/endpoints/controller/notifications.py

# -*- coding: utf-8 -*-
#
# Copyright (C) 2020 GNS3 Technologies Inc.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
API endpoints for controller notifications.
"""

import asyncio

from fastapi import APIRouter, WebSocket
from fastapi.responses import StreamingResponse
from starlette.endpoints import WebSocketEndpoint

from gns3server.controller import Controller

import logging
log = logging.getLogger(__name__)

router = APIRouter()


@router.get("")
async def http_notification():
    """
    Receive controller notifications about the controller from HTTP stream.
    """

    async def event_stream():

        with Controller.instance().notification.controller_queue() as queue:
            while True:
                msg = await queue.get_json(5)
                yield ("{}\n".format(msg)).encode("utf-8")

    return StreamingResponse(event_stream(), media_type="application/json")


@router.websocket_route("/ws")
class ControllerWebSocketNotifications(WebSocketEndpoint):
    """
    Receive controller notifications about the controller from WebSocket stream.
    """

    async def on_connect(self, websocket: WebSocket) -> None:

        await websocket.accept()
        log.info(f"New client {websocket.client.host}:{websocket.client.port} has connected to controller WebSocket")

        self._notification_task = asyncio.ensure_future(self._stream_notifications(websocket=websocket))

    async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:

        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 _stream_notifications(self, websocket: WebSocket) -> None:

        with Controller.instance().notifications.queue() as queue:
            while True:
                notification = await queue.get_json(5)
                await websocket.send_text(notification)