jumpserver/apps/notifications/ws.py

64 lines
1.9 KiB
Python
Raw Normal View History

2021-06-11 09:11:29 +00:00
import threading
import json
from channels.generic.websocket import JsonWebsocketConsumer
from common.utils import get_logger
from common.db.utils import safe_db_connection
2021-06-11 09:11:29 +00:00
from .site_msg import SiteMessageUtil
from .signal_handlers import new_site_msg_chan
2021-06-11 09:11:29 +00:00
logger = get_logger(__name__)
class SiteMsgWebsocket(JsonWebsocketConsumer):
refresh_every_seconds = 10
sub = None
2021-06-11 09:11:29 +00:00
def connect(self):
user = self.scope["user"]
if user.is_authenticated:
self.accept()
self.sub = self.watch_recv_new_site_msg()
2021-06-11 09:11:29 +00:00
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
2021-06-11 09:11:29 +00:00
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 self.sub:
self.sub.unsubscribe()