From 0170799be963b4c35116e725086fb9b2aecf291a Mon Sep 17 00:00:00 2001
From: Love98 <liaocai98@gmail.com>
Date: Fri, 27 Sep 2024 13:05:47 +0800
Subject: [PATCH] feat: support send alert with specific providers

---
 object/site_cache.go        |  4 ++--
 object/site_timer_health.go | 32 ++++++++++++++++++++------------
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/object/site_cache.go b/object/site_cache.go
index 5ac7ce9..6dd1aa5 100644
--- a/object/site_cache.go
+++ b/object/site_cache.go
@@ -107,8 +107,8 @@ func refreshSiteMap() error {
 		}
 
 		newSiteMap[strings.ToLower(site.Domain)] = site
-		if site.EnableAlert {
-			newHealthCheckNeededDomains = append(newHealthCheckNeededDomains, site.Domain)
+		if !shouldStopHealthCheck(site) {
+			newHealthCheckNeededDomains = append(newHealthCheckNeededDomains, strings.ToLower(site.Domain))
 		}
 		for _, domain := range site.OtherDomains {
 			if domain != "" {
diff --git a/object/site_timer_health.go b/object/site_timer_health.go
index a6bf210..98e17e5 100644
--- a/object/site_timer_health.go
+++ b/object/site_timer_health.go
@@ -16,6 +16,7 @@ package object
 
 import (
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/casdoor/casdoor-go-sdk/casdoorsdk"
@@ -58,17 +59,19 @@ func healthCheck(site *Site, domain string) error {
 	if user == nil {
 		return fmt.Errorf("user not found")
 	}
-	if user.Email != "" {
-		err = casdoorsdk.SendEmail("CasWAF HealthCheck Check Alert", pingResponse, "CasWAF", user.Email)
-	}
-	if err != nil {
-		fmt.Println(err)
-	}
-	if user.Phone != "" {
-		err = casdoorsdk.SendSms(pingResponse, user.Phone)
-	}
-	if err != nil {
-		fmt.Println(err)
+	for _, provider := range site.AlertProviders {
+		if strings.HasPrefix(provider, "Email/") {
+			err := casdoorsdk.SendEmailByProvider("CasWAF HealthCheck Check Alert", pingResponse, "CasWAF", provider[6:], user.Email)
+			if err != nil {
+				fmt.Println(err)
+			}
+		}
+		if strings.HasPrefix(provider, "SMS/") {
+			err := casdoorsdk.SendSmsByProvider(pingResponse, provider[4:], user.Phone)
+			if err != nil {
+				fmt.Println(err)
+			}
+		}
 	}
 	return nil
 }
@@ -83,7 +86,8 @@ func startHealthCheckLoop() {
 		go func() {
 			for {
 				site := GetSiteByDomain(domain)
-				if site == nil || !site.EnableAlert || site.Domain == "" || site.Status == "Inactive" {
+				if shouldStopHealthCheck(site) {
+					delete(healthCheckTryTimesMap, domain)
 					return
 				}
 
@@ -96,3 +100,7 @@ func startHealthCheckLoop() {
 		}()
 	}
 }
+
+func shouldStopHealthCheck(site *Site) bool {
+	return site == nil || !site.EnableAlert || site.Domain == "" || site.Status == "Inactive"
+}