openmedialibrary/oml/websocket.py

# -*- coding: utf-8 -*-


from tornado.websocket import WebSocketHandler
from tornado.ioloop import IOLoop
import json

from oxtornado import json_dumps

import state
import settings

import logging
logger = logging.getLogger(__name__)


class Handler(WebSocketHandler):

    def initialize(self, public=False):
        self._public = public

    def check_origin(self, origin):
        # allow access to websocket from site, installer and loader (local file)
        return self.request.host in origin or \
            origin in (
                'http://127.0.0.1:9841',
                'http://127.0.0.1:9842',
                'file://',
                'null'
            )

    def open(self):
        if self.request.headers['origin'] not in ('null', 'file://', 'http://127.0.0.1:9842') \
            and self.request.host not in self.request.headers['origin']:
            logger.debug('reject cross site attempt to open websocket %s', self.request)
            self.close()
        if self not in state.websockets:
            state.websockets.append(self)
        if state.update:
            trigger_event('updatestatus', state.update._status)
        else:
            from user.models import User
            trigger_event('status', {
                'id': settings.USER_ID,
                'online': state.online
            })
            for u in User.query.filter(User.id!=settings.USER_ID).filter_by(peered=True).all():
                u.trigger_status()

    #websocket calls
    def on_message(self, message):
        try:
            action, data = json.loads(message)
        except json.decoder.JSONDecodeError:
            logger.debug('invalid websocket message: %s', message)
            return
        if state.tasks:
            state.tasks.queue(action, data)

    def on_close(self):
        if self in state.websockets:
            state.websockets.remove(self)

    def post(self, event, data):
        message = json_dumps([event, data])
        if self.ws_connection is None:
            self.on_close()
        else:
            state.main.add_callback(lambda: self._write_message(message))

    async def _write_message(self, message):
        try:
            task = self.write_message(message)
            await task
        except tornado.iostream.StreamClosedError as e:
            self.on_close()
        except tornado.websocket.WebSocketClosedError as e:
            self.on_close()


def trigger_event(event, data):
    #if len(state.websockets):
    #    logger.debug('trigger event %s %s %s', event, data, len(state.websockets))
    for ws in state.websockets:
        try:
            ws.post(event, data)
        except:
            logger.debug('failed to send to ws %s %s %s', ws, event, data, exc_info=True)