refactor: 重构消息

pull/6824/head
xinwen 2021-09-13 17:16:20 +08:00 committed by Jiangjie.Bai
parent f64740c2db
commit d34c7edb00
3 changed files with 96 additions and 66 deletions

View File

@ -1,9 +1,12 @@
from typing import Iterable from typing import Iterable
import traceback import traceback
from itertools import chain from itertools import chain
import time
from celery import shared_task from celery import shared_task
from django.utils.translation import gettext_lazy as _
from common.utils.timezone import now
from common.utils import lazyproperty from common.utils import lazyproperty
from users.models import User from users.models import User
from notifications.backends import BACKEND from notifications.backends import BACKEND
@ -97,29 +100,50 @@ class Message(metaclass=MessageType):
def get_common_msg(self) -> dict: def get_common_msg(self) -> dict:
raise NotImplementedError raise NotImplementedError
def get_text_msg(self) -> dict:
return self.common_msg
def get_html_msg(self) -> dict:
return self.common_msg
@lazyproperty @lazyproperty
def common_msg(self) -> dict: def common_msg(self) -> dict:
return self.get_common_msg() return self.get_common_msg()
@lazyproperty
def text_msg(self) -> dict:
return self.get_text_msg()
@lazyproperty
def html_msg(self) -> dict:
return self.get_html_msg()
# -------------------------------------------------------------- # --------------------------------------------------------------
# 支持不同发送消息的方式定义自己的消息内容,比如有些支持 html 标签 # 支持不同发送消息的方式定义自己的消息内容,比如有些支持 html 标签
def get_dingtalk_msg(self) -> dict: def get_dingtalk_msg(self) -> dict:
return self.common_msg # 钉钉相同的消息一天只能发一次,所以给所有消息添加基于时间的序号,使他们不相同
message = self.text_msg['message']
suffix = _('\nTime: {}').format(now())
return {
'subject': self.text_msg['subject'],
'message': message + suffix
}
def get_wecom_msg(self) -> dict: def get_wecom_msg(self) -> dict:
return self.common_msg return self.text_msg
def get_feishu_msg(self) -> dict: def get_feishu_msg(self) -> dict:
return self.common_msg return self.text_msg
def get_email_msg(self) -> dict: def get_email_msg(self) -> dict:
return self.common_msg return self.html_msg
def get_site_msg_msg(self) -> dict: def get_site_msg_msg(self) -> dict:
return self.common_msg return self.html_msg
def get_sms_msg(self) -> dict: def get_sms_msg(self) -> dict:
return self.common_msg return self.text_msg
# -------------------------------------------------------------- # --------------------------------------------------------------

View File

@ -9,6 +9,7 @@ from notifications.notifications import SystemMessage
from terminal.models import Session, Command from terminal.models import Session, Command
from notifications.models import SystemMsgSubscription from notifications.models import SystemMsgSubscription
from notifications.backends import BACKEND from notifications.backends import BACKEND
from common.utils import lazyproperty
logger = get_logger(__name__) logger = get_logger(__name__)
@ -23,12 +24,8 @@ class CommandAlertMixin:
_get_message: Callable _get_message: Callable
message_type_label: str message_type_label: str
def get_dingtalk_msg(self) -> str: @lazyproperty
msg = self._get_message() def subject(self):
msg = msg.replace('<br>', '')
return msg
def get_subject(self):
_input = self.command['input'] _input = self.command['input']
if isinstance(_input, str): if isinstance(_input, str):
_input = _input.replace('\r\n', ' ').replace('\r', ' ').replace('\n', ' ') _input = _input.replace('\r\n', ' ').replace('\r', ' ').replace('\n', ' ')
@ -70,7 +67,35 @@ class CommandAlertMessage(CommandAlertMixin, SystemMessage):
def __init__(self, command): def __init__(self, command):
self.command = command self.command = command
def _get_message(self): def get_text_msg(self) -> dict:
command = self.command
session = Session.objects.get(id=command['session'])
session_detail_url = reverse(
'api-terminal:session-detail', kwargs={'pk': command['session']},
external=True, api_to_ui=True
)
message = _("""
Command: %(command)s
Asset: %(hostname)s (%(host_ip)s)
User: %(user)s
Level: %(risk_level)s
Session: %(session_detail_url)s?oid=%(oid)s
""") % {
'command': command['input'],
'hostname': command['asset'],
'host_ip': session.asset_obj.ip,
'user': command['user'],
'risk_level': Command.get_risk_level_str(command['risk_level']),
'session_detail_url': session_detail_url,
'oid': session.org_id
}
return {
'subject': self.subject,
'message': message
}
def get_html_msg(self) -> dict:
command = self.command command = self.command
session = Session.objects.get(id=command['session']) session = Session.objects.get(id=command['session'])
session_detail_url = reverse( session_detail_url = reverse(
@ -98,15 +123,9 @@ class CommandAlertMessage(CommandAlertMixin, SystemMessage):
'session_detail_url': session_detail_url, 'session_detail_url': session_detail_url,
'oid': session.org_id 'oid': session.org_id
} }
return message
def get_common_msg(self):
msg = self._get_message()
subject = self.get_subject()
return { return {
'subject': subject, 'subject': self.subject,
'message': msg 'message': message
} }
@ -118,7 +137,7 @@ class CommandExecutionAlert(CommandAlertMixin, SystemMessage):
def __init__(self, command): def __init__(self, command):
self.command = command self.command = command
def _get_message(self): def get_html_msg(self) -> dict:
command = self.command command = self.command
_input = command['input'] _input = command['input']
_input = _input.replace('\n', '<br>') _input = _input.replace('\n', '<br>')
@ -141,13 +160,31 @@ class CommandExecutionAlert(CommandAlertMixin, SystemMessage):
'user': command['user'], 'user': command['user'],
'risk_level': Command.get_risk_level_str(command['risk_level']) 'risk_level': Command.get_risk_level_str(command['risk_level'])
} }
return message
def get_common_msg(self):
subject = self.get_subject()
message = self._get_message()
return { return {
'subject': subject, 'subject': self.subject,
'message': message
}
def get_text_msg(self) -> dict:
command = self.command
_input = command['input']
assets = ', '.join([str(asset) for asset in command['assets']])
message = _("""
Assets: %(assets)s
User: %(user)s
Level: %(risk_level)s
Commands 👇 ------------
%(command)s
------------------------
""") % {
'command': _input,
'assets': assets,
'user': command['user'],
'risk_level': Command.get_risk_level_str(command['risk_level'])
}
return {
'subject': self.subject,
'message': message 'message': message
} }

View File

@ -6,38 +6,7 @@ from common.utils import reverse, get_request_ip_or_data, get_request_user_agent
from notifications.notifications import UserMessage from notifications.notifications import UserMessage
class BaseUserMessage(UserMessage): class ResetPasswordMsg(UserMessage):
def get_text_msg(self) -> dict:
raise NotImplementedError
def get_html_msg(self) -> dict:
raise NotImplementedError
@lazyproperty
def text_msg(self) -> dict:
return self.get_text_msg()
@lazyproperty
def html_msg(self) -> dict:
return self.get_html_msg()
def get_dingtalk_msg(self) -> dict:
return self.text_msg
def get_wecom_msg(self) -> dict:
return self.text_msg
def get_feishu_msg(self) -> dict:
return self.text_msg
def get_email_msg(self) -> dict:
return self.html_msg
def get_site_msg_msg(self) -> dict:
return self.html_msg
class ResetPasswordMsg(BaseUserMessage):
def get_text_msg(self) -> dict: def get_text_msg(self) -> dict:
user = self.user user = self.user
subject = _('Reset password') subject = _('Reset password')
@ -104,7 +73,7 @@ Login direct 👇
} }
class ResetPasswordSuccessMsg(BaseUserMessage): class ResetPasswordSuccessMsg(UserMessage):
def __init__(self, user, request): def __init__(self, user, request):
super().__init__(user) super().__init__(user)
self.ip_address = get_request_ip_or_data(request) self.ip_address = get_request_ip_or_data(request)
@ -187,7 +156,7 @@ Browser: %(browser)s
} }
class PasswordExpirationReminderMsg(BaseUserMessage): class PasswordExpirationReminderMsg(UserMessage):
def get_text_msg(self) -> dict: def get_text_msg(self) -> dict:
user = self.user user = self.user
@ -263,7 +232,7 @@ Login direct 👇
} }
class UserExpirationReminderMsg(BaseUserMessage): class UserExpirationReminderMsg(UserMessage):
def get_text_msg(self) -> dict: def get_text_msg(self) -> dict:
subject = _('Expiration notice') subject = _('Expiration notice')
message = _(""" message = _("""
@ -303,7 +272,7 @@ In order not to affect your normal work, please contact the administrator for co
} }
class ResetSSHKeyMsg(BaseUserMessage): class ResetSSHKeyMsg(UserMessage):
def get_text_msg(self) -> dict: def get_text_msg(self) -> dict:
subject = _('SSH Key Reset') subject = _('SSH Key Reset')
message = _(""" message = _("""
@ -347,7 +316,7 @@ Login direct 👇
} }
class ResetMFAMsg(BaseUserMessage): class ResetMFAMsg(UserMessage):
def get_text_msg(self) -> dict: def get_text_msg(self) -> dict:
subject = _('MFA Reset') subject = _('MFA Reset')
message = _(""" message = _("""