diff --git a/backend/constants/notification.go b/backend/constants/notification.go index cad3b19b..cf3da062 100644 --- a/backend/constants/notification.go +++ b/backend/constants/notification.go @@ -5,3 +5,9 @@ const ( NotificationTriggerOnTaskError = "notification_trigger_on_task_error" NotificationTriggerNever = "notification_trigger_never" ) + +const ( + NotificationTypeMail = "notification_type_mail" + NotificationTypeDingTalk = "notification_type_ding_talk" + NotificationTypeWechat = "notification_type_wechat" +) diff --git a/backend/model/user.go b/backend/model/user.go index e37783c9..9dadec0f 100644 --- a/backend/model/user.go +++ b/backend/model/user.go @@ -24,9 +24,10 @@ type User struct { } type UserSetting struct { - NotificationTrigger string `json:"notification_trigger" bson:"notification_trigger"` - DingTalkRobotWebhook string `json:"ding_talk_robot_webhook" bson:"ding_talk_robot_webhook"` - WechatRobotWebhook string `json:"wechat_robot_webhook" bson:"wechat_robot_webhook"` + NotificationTrigger string `json:"notification_trigger" bson:"notification_trigger"` + DingTalkRobotWebhook string `json:"ding_talk_robot_webhook" bson:"ding_talk_robot_webhook"` + WechatRobotWebhook string `json:"wechat_robot_webhook" bson:"wechat_robot_webhook"` + EnabledNotifications []string `json:"enabled_notifications" bson:"enabled_notifications"` } func (user *User) Save() error { diff --git a/backend/routes/user.go b/backend/routes/user.go index 9408024a..fcca967e 100644 --- a/backend/routes/user.go +++ b/backend/routes/user.go @@ -96,16 +96,7 @@ func PutUser(c *gin.Context) { } // 添加用户 - user := model.User{ - Username: strings.ToLower(reqData.Username), - Password: utils.EncryptPassword(reqData.Password), - Role: reqData.Role, - Email: reqData.Email, - Setting: model.UserSetting{ - NotificationTrigger: constants.NotificationTriggerNever, - }, - } - if err := user.Add(); err != nil { + if err := services.CreateNewUser(reqData.Username, reqData.Password, reqData.Role, reqData.Email); err != nil { HandleError(http.StatusInternalServerError, c, err) return } @@ -238,6 +229,7 @@ func PostMe(c *gin.Context) { if reqBody.Setting.WechatRobotWebhook != "" { user.Setting.WechatRobotWebhook = reqBody.Setting.WechatRobotWebhook } + user.Setting.EnabledNotifications = reqBody.Setting.EnabledNotifications if err := user.Save(); err != nil { HandleError(http.StatusInternalServerError, c, err) return diff --git a/backend/services/task.go b/backend/services/task.go index 507fb8e6..da4eabc8 100644 --- a/backend/services/task.go +++ b/backend/services/task.go @@ -490,23 +490,7 @@ func ExecuteTask(id int) { // 如果发生错误，则发送通知 t, _ = model.GetTask(t.Id) if user.Setting.NotificationTrigger == constants.NotificationTriggerOnTaskEnd || user.Setting.NotificationTrigger == constants.NotificationTriggerOnTaskError { - if user.Email != "" { - go func() { - SendTaskEmail(user, t, spider) - }() - } - - if user.Setting.DingTalkRobotWebhook != "" { - go func() { - SendTaskDingTalk(user, t, spider) - }() - } - - if user.Setting.WechatRobotWebhook != "" { - go func() { - SendTaskWechat(user, t, spider) - }() - } + SendNotifications(user, t, spider) } return } @@ -532,27 +516,7 @@ func ExecuteTask(id int) { // 如果是任务结束时发送通知，则发送通知 if user.Setting.NotificationTrigger == constants.NotificationTriggerOnTaskEnd { - if user.Email != "" { - go func() { - SendTaskEmail(user, t, spider) - }() - } - - if user.Email != "" { - go func() { - if user.Setting.DingTalkRobotWebhook != "" { - SendTaskDingTalk(user, t, spider) - } - }() - } - - if user.Email != "" { - go func() { - if user.Setting.WechatRobotWebhook != "" { - SendTaskWechat(user, t, spider) - } - }() - } + SendNotifications(user, t, spider) } // 保存任务 @@ -779,7 +743,7 @@ func GetTaskMarkdownContent(t model.Task, s model.Spider) string { errLog := "-" statusMsg := fmt.Sprintf(`%s`, t.Status) if t.Status == constants.StatusError { - errMsg = `（有错误）` + errMsg = `（有错误）` errLog = fmt.Sprintf(`%s`, t.Error) statusMsg = fmt.Sprintf(`%s`, t.Status) } @@ -859,6 +823,26 @@ func SendTaskWechat(u model.User, t model.Task, s model.Spider) { } } +func SendNotifications(u model.User, t model.Task, s model.Spider) { + if u.Email != "" && utils.StringArrayContains(u.Setting.EnabledNotifications, constants.NotificationTypeMail) { + go func() { + SendTaskEmail(u, t, s) + }() + } + + if u.Setting.DingTalkRobotWebhook != "" && utils.StringArrayContains(u.Setting.EnabledNotifications, constants.NotificationTypeDingTalk) { + go func() { + SendTaskDingTalk(u, t, s) + }() + } + + if u.Setting.WechatRobotWebhook != "" && utils.StringArrayContains(u.Setting.EnabledNotifications, constants.NotificationTypeWechat) { + go func() { + SendTaskWechat(u, t, s) + }() + } +} + func InitTaskExecutor() error { c := cron.New(cron.WithSeconds()) Exec = &Executor{ diff --git a/backend/services/user.go b/backend/services/user.go index 483a7408..a01e721b 100644 --- a/backend/services/user.go +++ b/backend/services/user.go @@ -9,21 +9,15 @@ import ( "github.com/gin-gonic/gin" "github.com/globalsign/mgo/bson" "github.com/spf13/viper" + "strings" "time" ) func InitUserService() error { - adminUser := model.User{ - Username: "admin", - Password: utils.EncryptPassword("admin"), - Role: constants.RoleAdmin, - Setting: model.UserSetting{ - NotificationTrigger: constants.NotificationTriggerNever, - }, - } - _ = adminUser.Add() + _ = CreateNewUser("admin", "admin", constants.RoleAdmin, "") return nil } + func MakeToken(user *model.User) (tokenStr string, err error) { token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{ "id": user.Id, @@ -96,6 +90,27 @@ func CheckToken(tokenStr string) (user model.User, err error) { return } +func CreateNewUser(username string, password string, role string, email string) error { + user := model.User{ + Username: strings.ToLower(username), + Password: utils.EncryptPassword(password), + Role: role, + Email: email, + Setting: model.UserSetting{ + NotificationTrigger: constants.NotificationTriggerNever, + EnabledNotifications: []string{ + constants.NotificationTypeMail, + constants.NotificationTypeDingTalk, + constants.NotificationTypeWechat, + }, + }, + } + if err := user.Add(); err != nil { + return err + } + return nil +} + func GetCurrentUser(c *gin.Context) *model.User { data, _ := c.Get("currentUser") return data.(*model.User) diff --git a/backend/utils/array.go b/backend/utils/array.go new file mode 100644 index 00000000..889430ed --- /dev/null +++ b/backend/utils/array.go @@ -0,0 +1,10 @@ +package utils + +func StringArrayContains(arr []string, str string) bool { + for _, s := range arr { + if s == str { + return true + } + } + return false +} diff --git a/frontend/src/i18n/zh.js b/frontend/src/i18n/zh.js index a3ea2768..136dbf3f 100644 --- a/frontend/src/i18n/zh.js +++ b/frontend/src/i18n/zh.js @@ -349,11 +349,12 @@ export default { 'Optional': '可选', // 设置 - 'Notification Trigger': '通知触发', + 'Notification Trigger Timing': '消息通知触发时机', 'On Task End': '当任务结束', 'On Task Error': '当任务发生错误', 'Never': '从不', 'DingTalk Robot Webhook': '钉钉机器人 Webhook', + 'Wechat Robot Webhook': '微信机器人 Webhook', // 其他 tagsView: { diff --git a/frontend/src/views/setting/Setting.vue b/frontend/src/views/setting/Setting.vue index 33a134a6..75dc856f 100644 --- a/frontend/src/views/setting/Setting.vue +++ b/frontend/src/views/setting/Setting.vue @@ -8,10 +8,8 @@ - - - - +
+ {{$t('On Task End')}} @@ -24,8 +22,22 @@ + + + {{$t('邮件')}} + {{$t('钉钉')}} + {{$t('企业微信')}} + + + + + - + + + +
@@ -64,12 +76,21 @@ export default { callback() } } + const validateWechatRobotWebhook = (rule, value, callback) => { + if (!value) return callback() + if (!value.match(/^https:\/\/qyapi.weixin.qq.com\/cgi-bin\/webhook\/send\?key=.+/i)) { + callback(new Error(this.$t('DingTalk Robot Webhook format invalid'))) + } else { + callback() + } + } return { - userInfo: { setting: {} }, + userInfo: { setting: { enabled_notifications: [] } }, rules: { password: [{ trigger: 'blur', validator: validatePass }], email: [{ trigger: 'blur', validator: validateEmail }], - 'setting.ding_talk_robot_webhook': [{ trigger: 'blur', validator: validateDingTalkRobotWebhook }] + 'setting.ding_talk_robot_webhook': [{ trigger: 'blur', validator: validateDingTalkRobotWebhook }], + 'setting.wechat_robot_webhook': [{ trigger: 'blur', validator: validateWechatRobotWebhook }] }, isShowDingTalkAppSecret: false } @@ -80,16 +101,12 @@ export default { if (!data) return {} this.userInfo = JSON.parse(data) if (!this.userInfo.setting) this.userInfo.setting = {} + if (!this.userInfo.setting.enabled_notifications) this.userInfo.setting.enabled_notifications = [] }, saveUserInfo () { this.$refs['setting-form'].validate(async valid => { if (!valid) return - const res = await this.$store.dispatch('user/postInfo', { - password: this.userInfo.password, - email: this.userInfo.email, - notification_trigger: this.userInfo.setting.notification_trigger, - ding_talk_robot_webhook: this.userInfo.setting.ding_talk_robot_webhook - }) + const res = await this.$store.dispatch('user/postInfo', this.userInfo) if (!res || res.error) { this.$message.error(res.error) } else {