crafty-4/app/classes/web/websocket_handler.py

import json
import logging
import asyncio
import sys

from urllib.parse import parse_qsl
from app.classes.models.users import Users
from app.classes.shared.helpers import helper
from app.classes.web.websocket_helper import websocket_helper
from app.classes.shared.console import console

logger = logging.getLogger(__name__)

try:
    import tornado.websocket

except ModuleNotFoundError as e:
    logger.critical("Import Error: Unable to load {} module".format(e, e.name))
    console.critical("Import Error: Unable to load {} module".format(e, e.name))
    sys.exit(1)

class SocketHandler(tornado.websocket.WebSocketHandler):

    def initialize(self, controller=None, tasks_manager=None, translator=None):
        self.controller = controller
        self.tasks_manager = tasks_manager
        self.translator = translator
        self.io_loop = tornado.ioloop.IOLoop.current()

    def get_remote_ip(self):
        remote_ip = self.request.headers.get("X-Real-IP") or \
                    self.request.headers.get("X-Forwarded-For") or \
                    self.request.remote_ip
        return remote_ip

    def get_user_id(self):
        user_data_cookie_raw = self.get_secure_cookie('user_data')

        if user_data_cookie_raw and user_data_cookie_raw.decode('utf-8'):
            user_data_cookie = user_data_cookie_raw.decode('utf-8')
            user_id = json.loads(user_data_cookie)['user_id']
            return user_id

    def check_auth(self):
        user_data_cookie_raw = self.get_secure_cookie('user_data')

        if user_data_cookie_raw and user_data_cookie_raw.decode('utf-8'):
            user_data_cookie = user_data_cookie_raw.decode('utf-8')
            user_id = json.loads(user_data_cookie)['user_id']
            query = Users.select().where(Users.user_id == user_id)
            if query.exists():
                return True
        return False


    def open(self):
        logger.debug('Checking WebSocket authentication')
        if self.check_auth():
            self.handle()
        else:
            websocket_helper.send_message(self, 'notification', 'Not authenticated for WebSocket connection')
            self.close()
            self.controller.management.add_to_audit_log_raw('unknown', 0, 0, 'Someone tried to connect via WebSocket without proper authentication', self.get_remote_ip())
            websocket_helper.broadcast('notification', 'Someone tried to connect via WebSocket without proper authentication')
            logger.warning('Someone tried to connect via WebSocket without proper authentication')

    def handle(self):
        self.page = self.get_query_argument('page')
        self.page_query_params = dict(parse_qsl(helper.remove_prefix(
            self.get_query_argument('page_query_params'),
            '?'
        )))
        websocket_helper.add_client(self)
        logger.debug('Opened WebSocket connection')
        # websocket_helper.broadcast('notification', 'New client connected')

    def on_message(self, rawMessage):

        logger.debug('Got message from WebSocket connection {}'.format(rawMessage))
        message = json.loads(rawMessage)
        logger.debug('Event Type: {}, Data: {}'.format(message['event'], message['data']))

    def on_close(self):
        websocket_helper.remove_client(self)
        logger.debug('Closed WebSocket connection')
        # websocket_helper.broadcast('notification', 'Client disconnected')

    async def write_message_int(self, message):
        self.write_message(message)
    
    def write_message_helper(self, message):
        asyncio.run_coroutine_threadsafe(self.write_message_int(message), self.io_loop.asyncio_loop)