jumpserver/apps/notifications/ws.py

import json
import time
from threading import Thread

from channels.generic.websocket import JsonWebsocketConsumer
from django.conf import settings

from common.db.utils import safe_db_connection
from common.sessions.cache import user_session_manager
from common.utils import get_logger
from .signal_handlers import new_site_msg_chan
from .site_msg import SiteMessageUtil

logger = get_logger(__name__)


class SiteMsgWebsocket(JsonWebsocketConsumer):
    sub = None
    refresh_every_seconds = 10

    @property
    def session(self):
        return self.scope['session']

    def connect(self):
        user = self.scope["user"]
        if user.is_authenticated:
            self.accept()
            user_session_manager.add_or_increment(self.session.session_key)
            self.sub = self.watch_recv_new_site_msg()
        else:
            self.close()

    def receive(self, text_data=None, bytes_data=None, **kwargs):
        data = json.loads(text_data)
        refresh_every_seconds = data.get('refresh_every_seconds')

        try:
            refresh_every_seconds = int(refresh_every_seconds)
        except Exception as e:
            logger.error(e)
            return

        if refresh_every_seconds > 0:
            self.refresh_every_seconds = refresh_every_seconds

    def send_unread_msg_count(self):
        user_id = self.scope["user"].id
        unread_count = SiteMessageUtil.get_user_unread_msgs_count(user_id)
        logger.debug('Send unread count to user: {} {}'.format(user_id, unread_count))
        self.send_json({'type': 'unread_count', 'unread_count': unread_count})

    def watch_recv_new_site_msg(self):
        ws = self
        user_id = str(self.scope["user"].id)

        # 先发一个消息再说
        with safe_db_connection():
            self.send_unread_msg_count()

        def handle_new_site_msg_recv(msg):
            users = msg.get('users', [])
            logger.debug('New site msg recv, message users: {}'.format(users))
            if user_id in users:
                ws.send_unread_msg_count()

        return new_site_msg_chan.subscribe(handle_new_site_msg_recv)

    def disconnect(self, code):
        if not self.sub:
            return
        self.sub.unsubscribe()

        user_session_manager.decrement_or_remove(self.session.session_key)
        if self.should_delete_session():
            thread = Thread(target=self.delay_delete_session)
            thread.start()

    def should_delete_session(self):
        return (self.session.modified or settings.SESSION_SAVE_EVERY_REQUEST) and \
            not self.session.is_empty() and \
            self.session.get_expire_at_browser_close() and \
            not user_session_manager.check_active(self.session.session_key)

    def delay_delete_session(self):
        timeout = 3
        check_interval = 0.5

        start_time = time.time()
        while time.time() - start_time < timeout:
            time.sleep(check_interval)
            if user_session_manager.check_active(self.session.session_key):
                return

        self.delete_session()

    def delete_session(self):
        try:
            self.session.delete()
        except Exception as e:
            logger.info(f'delete session error: {e}')