From d34c7edb0018dc39fb1398bd5d29b240618eb83c Mon Sep 17 00:00:00 2001 From: xinwen Date: Mon, 13 Sep 2021 17:16:20 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84=E6=B6=88?= =?UTF-8?q?=E6=81=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/notifications/notifications.py | 36 ++++++++++--- apps/terminal/notifications.py | 83 +++++++++++++++++++++-------- apps/users/notifications.py | 43 +++------------ 3 files changed, 96 insertions(+), 66 deletions(-) diff --git a/apps/notifications/notifications.py b/apps/notifications/notifications.py index 97ff30469..ab2f25492 100644 --- a/apps/notifications/notifications.py +++ b/apps/notifications/notifications.py @@ -1,9 +1,12 @@ from typing import Iterable import traceback from itertools import chain +import time 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 users.models import User from notifications.backends import BACKEND @@ -97,29 +100,50 @@ class Message(metaclass=MessageType): def get_common_msg(self) -> dict: raise NotImplementedError + def get_text_msg(self) -> dict: + return self.common_msg + + def get_html_msg(self) -> dict: + return self.common_msg + @lazyproperty def common_msg(self) -> dict: 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 标签 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: - return self.common_msg + return self.text_msg def get_feishu_msg(self) -> dict: - return self.common_msg + return self.text_msg def get_email_msg(self) -> dict: - return self.common_msg + return self.html_msg def get_site_msg_msg(self) -> dict: - return self.common_msg + return self.html_msg def get_sms_msg(self) -> dict: - return self.common_msg + return self.text_msg # -------------------------------------------------------------- diff --git a/apps/terminal/notifications.py b/apps/terminal/notifications.py index 22e1acf6d..d7581d1e1 100644 --- a/apps/terminal/notifications.py +++ b/apps/terminal/notifications.py @@ -9,6 +9,7 @@ from notifications.notifications import SystemMessage from terminal.models import Session, Command from notifications.models import SystemMsgSubscription from notifications.backends import BACKEND +from common.utils import lazyproperty logger = get_logger(__name__) @@ -23,12 +24,8 @@ class CommandAlertMixin: _get_message: Callable message_type_label: str - def get_dingtalk_msg(self) -> str: - msg = self._get_message() - msg = msg.replace('
', '') - return msg - - def get_subject(self): + @lazyproperty + def subject(self): _input = self.command['input'] if isinstance(_input, str): _input = _input.replace('\r\n', ' ').replace('\r', ' ').replace('\n', ' ') @@ -70,7 +67,35 @@ class CommandAlertMessage(CommandAlertMixin, SystemMessage): def __init__(self, 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 session = Session.objects.get(id=command['session']) session_detail_url = reverse( @@ -98,15 +123,9 @@ class CommandAlertMessage(CommandAlertMixin, SystemMessage): 'session_detail_url': session_detail_url, 'oid': session.org_id } - return message - - def get_common_msg(self): - msg = self._get_message() - subject = self.get_subject() - return { - 'subject': subject, - 'message': msg + 'subject': self.subject, + 'message': message } @@ -118,7 +137,7 @@ class CommandExecutionAlert(CommandAlertMixin, SystemMessage): def __init__(self, command): self.command = command - def _get_message(self): + def get_html_msg(self) -> dict: command = self.command _input = command['input'] _input = _input.replace('\n', '
') @@ -141,13 +160,31 @@ class CommandExecutionAlert(CommandAlertMixin, SystemMessage): 'user': command['user'], '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 { - '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 } diff --git a/apps/users/notifications.py b/apps/users/notifications.py index 7d860ca14..a129c624e 100644 --- a/apps/users/notifications.py +++ b/apps/users/notifications.py @@ -6,38 +6,7 @@ from common.utils import reverse, get_request_ip_or_data, get_request_user_agent from notifications.notifications import UserMessage -class BaseUserMessage(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): +class ResetPasswordMsg(UserMessage): def get_text_msg(self) -> dict: user = self.user subject = _('Reset password') @@ -104,7 +73,7 @@ Login direct 👇 } -class ResetPasswordSuccessMsg(BaseUserMessage): +class ResetPasswordSuccessMsg(UserMessage): def __init__(self, user, request): super().__init__(user) 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: user = self.user @@ -263,7 +232,7 @@ Login direct 👇 } -class UserExpirationReminderMsg(BaseUserMessage): +class UserExpirationReminderMsg(UserMessage): def get_text_msg(self) -> dict: subject = _('Expiration notice') 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: subject = _('SSH Key Reset') message = _(""" @@ -347,7 +316,7 @@ Login direct 👇 } -class ResetMFAMsg(BaseUserMessage): +class ResetMFAMsg(UserMessage): def get_text_msg(self) -> dict: subject = _('MFA Reset') message = _("""