jumpserver/apps/notifications/signal_handlers.py

91 lines
2.8 KiB
Python

import inspect
from importlib import import_module
from django.apps import AppConfig
from django.db.models.signals import post_migrate
from django.db.models.signals import post_save
from django.dispatch import receiver
from django.utils.functional import LazyObject
from common.decorators import on_transaction_commit
from common.utils import get_logger
from common.utils.connection import RedisPubSub
from notifications.backends import BACKEND
from users.models import User
from .models import MessageContent, SystemMsgSubscription, UserMsgSubscription
from .notifications import SystemMessage
logger = get_logger(__name__)
class NewSiteMsgSubPub(LazyObject):
def _setup(self):
self._wrapped = RedisPubSub('notifications.SiteMessageCome')
new_site_msg_chan = NewSiteMsgSubPub()
@receiver(post_save, sender=MessageContent)
@on_transaction_commit
def on_site_message_create(sender, instance, created, **kwargs):
if not created:
return
logger.debug('New site msg created, publish it')
user_ids = instance.users.all().values_list('id', flat=True)
user_ids = [str(i) for i in user_ids]
data = {
'id': str(instance.id),
'subject': instance.subject,
'message': instance.message,
'users': user_ids
}
new_site_msg_chan.publish(data)
@receiver(post_migrate, dispatch_uid='notifications.signal_handlers.create_system_messages')
def create_system_messages(app_config: AppConfig, **kwargs):
try:
notifications_module = import_module('.notifications', app_config.module.__package__)
for name, obj in notifications_module.__dict__.items():
if name.startswith('_'):
continue
if not inspect.isclass(obj):
continue
if not issubclass(obj, SystemMessage):
continue
attrs = obj.__dict__
if 'message_type_label' not in attrs:
continue
if 'category' not in attrs:
continue
if 'category_label' not in attrs:
continue
message_type = obj.get_message_type()
sub, created = SystemMsgSubscription.objects.get_or_create(message_type=message_type)
if created:
obj.post_insert_to_db(sub)
logger.info(
f'Create SystemMsgSubscription: package={app_config.module.__package__} type={message_type}')
except ModuleNotFoundError:
pass
@receiver(post_save, sender=User)
def on_user_post_save(sender, instance, created, **kwargs):
if not created:
return
receive_backends = []
# Todo: IDE 识别不了 get_account
for backend in BACKEND:
if backend.get_account(instance):
receive_backends.append(backend)
UserMsgSubscription.objects.create(user=instance, receive_backends=receive_backends)