diff --git a/spug_api/apps/notify/models.py b/spug_api/apps/notify/models.py index c7f122c..c5c7516 100644 --- a/spug_api/apps/notify/models.py +++ b/spug_api/apps/notify/models.py @@ -4,6 +4,7 @@ from django.db import models from django.core.cache import cache from libs import ModelMixin, human_datetime +from libs.channel import Channel import time @@ -31,6 +32,7 @@ class Notify(models.Model, ModelMixin): if not with_quiet or time.time() - cache.get('spug:notify_quiet', 0) > 3600: cache.set('spug:notify_quiet', time.time()) cls.objects.create(source=source, title=title, type=type, content=content) + Channel.send_notify(title, content) def __repr__(self): return '' % self.title diff --git a/spug_api/consumer/consumers.py b/spug_api/consumer/consumers.py index 92cbfd2..b4d32fd 100644 --- a/spug_api/consumer/consumers.py +++ b/spug_api/consumer/consumers.py @@ -3,6 +3,7 @@ # Released under the AGPL-3.0 License. from channels.generic.websocket import WebsocketConsumer from django_redis import get_redis_connection +from asgiref.sync import async_to_sync from apps.host.models import Host from threading import Thread import json @@ -88,3 +89,18 @@ class SSHConsumer(WebsocketConsumer): self.chan = self.ssh.invoke_shell(term='xterm') self.chan.transport.set_keepalive(30) Thread(target=self.loop_read).start() + + +class NotifyConsumer(WebsocketConsumer): + def connect(self): + async_to_sync(self.channel_layer.group_add)('notify', self.channel_name) + self.accept() + + def disconnect(self, code): + async_to_sync(self.channel_layer.group_discard)('notify', self.channel_name) + + def receive(self, **kwargs): + self.send(text_data='pong') + + def notify_message(self, event): + self.send(text_data=json.dumps(event)) diff --git a/spug_api/consumer/routing.py b/spug_api/consumer/routing.py index 07623e7..2b9aeff 100644 --- a/spug_api/consumer/routing.py +++ b/spug_api/consumer/routing.py @@ -10,5 +10,6 @@ ws_router = AuthMiddleware( URLRouter([ path('ws/exec//', ExecConsumer), path('ws/ssh//', SSHConsumer), + path('ws/notify/', NotifyConsumer), ]) ) diff --git a/spug_api/libs/channel.py b/spug_api/libs/channel.py index d04382a..5b66a95 100644 --- a/spug_api/libs/channel.py +++ b/spug_api/libs/channel.py @@ -25,3 +25,12 @@ class Channel: 'pkey': pkey } async_to_sync(layer.send)('ssh_exec', message) + + @staticmethod + def send_notify(title, content): + message = { + 'type': 'notify.message', + 'title': title, + 'content': content + } + async_to_sync(layer.group_send)('notify', message) diff --git a/spug_api/spug/routing.py b/spug_api/spug/routing.py index 76314f1..41caacf 100644 --- a/spug_api/spug/routing.py +++ b/spug_api/spug/routing.py @@ -2,11 +2,12 @@ # Copyright: (c) # Released under the AGPL-3.0 License. from channels.routing import ProtocolTypeRouter, ChannelNameRouter -from consumer import routing, executors +from consumer import routing, executors, consumers application = ProtocolTypeRouter({ 'channel': ChannelNameRouter({ 'ssh_exec': executors.SSHExecutor, + 'notify_message': consumers.NotifyConsumer, }), 'websocket': routing.ws_router }) diff --git a/spug_web/src/layout/Notification.js b/spug_web/src/layout/Notification.js index 13eed2d..28e28f8 100644 --- a/spug_web/src/layout/Notification.js +++ b/spug_web/src/layout/Notification.js @@ -1,11 +1,12 @@ import React, { useState, useEffect } from 'react'; -import { Menu, List, Dropdown, Badge } from 'antd'; +import { Menu, List, Dropdown, Badge, Button, notification } from 'antd'; import { CheckOutlined, NotificationOutlined } from '@ant-design/icons'; -import { http } from 'libs'; +import { http, X_TOKEN } from 'libs'; import moment from 'moment'; import styles from './layout.module.less'; -let interval; +let ws = {readyState: 3}; +let timer; export default function () { const [loading, setLoading] = useState(false); @@ -14,10 +15,19 @@ export default function () { useEffect(() => { fetch(); - interval = setInterval(fetch, 60000); + listen(); + timer = setInterval(() => { + if (ws.readyState === 1) { + ws.send('ping') + } else if (ws.readyState === 3) { + listen() + } + }, 10000) return () => { - if (interval) clearInterval(interval) + if (timer) clearInterval(timer); + if (ws.close) ws.close() } + // eslint-disable-next-line react-hooks/exhaustive-deps }, []) function fetch() { @@ -30,6 +40,24 @@ export default function () { .finally(() => setLoading(false)) } + function listen() { + const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; + ws = new WebSocket(`${protocol}//${window.location.host}/api/ws/notify/?x-token=${X_TOKEN}`); + ws.onopen = () => ws.send('ok'); + ws.onmessage = e => { + if (e.data === 'pong') { + } else { + fetch(); + const {title, content} = JSON.parse(e.data); + const key = `open${Date.now()}`; + const btn = ( + + ); + notification.warning({message: title, description: content, btn, key, top: 64, duration: null}) + } + } + } + function handleRead(e, item) { e.stopPropagation(); if (reads.indexOf(item.id) === -1) {