🔊 v0.3.0 计划任务（定期备份等场景）

cmd/dashboard/controller/controller.go

@@ -22,7 +22,7 @@ func ServeWeb(port uint) {
 	r.Use(mygin.RecordPath)
 	r.SetFuncMap(template.FuncMap{
 		"tf": func(t time.Time) string {
-			return t.Format("2006年1月2号")
+			return t.Format("2006年1月2号 15:04:05")
 		},
 		"safe": func(s string) template.HTML {
 			return template.HTML(s)

cmd/dashboard/controller/member_api.go

@@ -9,10 +9,12 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
-	"github.com/naiba/com"
+	"github.com/robfig/cron/v3"
 
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/pkg/mygin"
+	"github.com/naiba/nezha/pkg/utils"
+	pb "github.com/naiba/nezha/proto"
 	"github.com/naiba/nezha/service/alertmanager"
 	"github.com/naiba/nezha/service/dao"
 )
@@ -34,6 +36,7 @@ func (ma *memberAPI) serve() {
 	mr.POST("/logout", ma.logout)
 	mr.POST("/server", ma.addOrEditServer)
 	mr.POST("/monitor", ma.addOrEditMonitor)
+	mr.POST("/cron", ma.addOrEditCron)
 	mr.POST("/notification", ma.addOrEditNotification)
 	mr.POST("/alert-rule", ma.addOrEditAlertRule)
 	mr.POST("/setting", ma.updateSetting)
@@ -70,6 +73,17 @@ func (ma *memberAPI) delete(c *gin.Context) {
 		if err == nil {
 			err = dao.DB.Delete(&model.MonitorHistory{}, "monitor_id = ?", id).Error
 		}
+	case "cron":
+
+		err = dao.DB.Delete(&model.Cron{}, "id = ?", id).Error
+		if err == nil {
+			dao.CronLock.RLock()
+			defer dao.CronLock.RUnlock()
+			if dao.Crons[id].CronID != 0 {
+				dao.Cron.Remove(dao.Crons[id].CronID)
+			}
+			delete(dao.Crons, id)
+		}
 	case "alert-rule":
 		err = dao.DB.Delete(&model.AlertRule{}, "id = ?", id).Error
 		if err == nil {
@@ -109,7 +123,7 @@ func (ma *memberAPI) addOrEditServer(c *gin.Context) {
 		s.ID = sf.ID
 		s.Tag = sf.Tag
 		if sf.ID == 0 {
-			s.Secret = com.MD5(fmt.Sprintf("%s%s%d", time.Now(), sf.Name, admin.ID))
+			s.Secret = utils.MD5(fmt.Sprintf("%s%s%d", time.Now(), sf.Name, admin.ID))
 			s.Secret = s.Secret[:10]
 			err = dao.DB.Create(&s).Error
 		} else {
@@ -179,6 +193,72 @@ func (ma *memberAPI) addOrEditMonitor(c *gin.Context) {
 	})
 }
 
+type cronForm struct {
+	ID             uint64
+	Name           string
+	Scheduler      string
+	Command        string
+	ServersRaw     string
+	PushSuccessful string
+}
+
+func (ma *memberAPI) addOrEditCron(c *gin.Context) {
+	var cf cronForm
+	var cr model.Cron
+	err := c.ShouldBindJSON(&cf)
+	if err == nil {
+		cr.Name = cf.Name
+		cr.Scheduler = cf.Scheduler
+		cr.Command = cf.Command
+		cr.ServersRaw = cf.ServersRaw
+		cr.PushSuccessful = cf.PushSuccessful == "on"
+		cr.ID = cf.ID
+		err = json.Unmarshal([]byte(cf.ServersRaw), &cr.Servers)
+	}
+	if err == nil {
+		_, err = cron.ParseStandard(cr.Scheduler)
+	}
+	if err == nil {
+		if cf.ID == 0 {
+			err = dao.DB.Create(&cr).Error
+		} else {
+			err = dao.DB.Save(&cr).Error
+		}
+	}
+
+	if err != nil {
+		c.JSON(http.StatusOK, model.Response{
+			Code:    http.StatusBadRequest,
+			Message: fmt.Sprintf("请求错误：%s", err),
+		})
+		return
+	}
+
+	if cr.CronID != 0 {
+		dao.Cron.Remove(cr.CronID)
+	}
+
+	cr.CronID, err = dao.Cron.AddFunc(cr.Scheduler, func() {
+		dao.ServerLock.RLock()
+		defer dao.ServerLock.RUnlock()
+		for j := 0; j < len(cr.Servers); j++ {
+			if dao.ServerList[cr.Servers[j]].TaskStream != nil {
+				dao.ServerList[cr.Servers[j]].TaskStream.Send(&pb.Task{
+					Id:   cr.ID,
+					Data: cr.Command,
+					Type: model.TaskTypeCommand,
+				})
+			} else {
+				alertmanager.SendNotification(fmt.Sprintf("计划任务：%s，服务器：%d 离线，无法执行。", cr.Name, cr.Servers[j]))
+			}
+		}
+	})
+
+	c.JSON(http.StatusOK, model.Response{
+		Code: http.StatusOK,
+	})
+}
+
 type notificationForm struct {
 	ID            uint64
 	Name          string

cmd/dashboard/controller/member_page.go

@@ -24,6 +24,7 @@ func (mp *memberPage) serve() {
 	}))
 	mr.GET("/server", mp.server)
 	mr.GET("/monitor", mp.monitor)
+	mr.GET("/cron", mp.cron)
 	mr.GET("/notification", mp.notification)
 	mr.GET("/setting", mp.setting)
 }
@@ -46,6 +47,15 @@ func (mp *memberPage) monitor(c *gin.Context) {
 	}))
 }
 
+func (mp *memberPage) cron(c *gin.Context) {
+	var crons []model.Cron
+	dao.DB.Find(&crons)
+	c.HTML(http.StatusOK, "dashboard/cron", mygin.CommonEnvironment(c, gin.H{
+		"Title": "计划任务",
+		"Crons": crons,
+	}))
+}
+
 func (mp *memberPage) notification(c *gin.Context) {
 	var nf []model.Notification
 	dao.DB.Find(&nf)

cmd/dashboard/controller/oauth2.go

@@ -6,14 +6,13 @@ import (
 	"net/http"
 	"strings"
 
-	"github.com/naiba/com"
-
 	"github.com/gin-gonic/gin"
 	GitHubAPI "github.com/google/go-github/github"
 	"golang.org/x/oauth2"
 
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/pkg/mygin"
+	"github.com/naiba/nezha/pkg/utils"
 	"github.com/naiba/nezha/service/dao"
 )
 
@@ -28,7 +27,7 @@ func (oa *oauth2controller) serve() {
 }
 
 func (oa *oauth2controller) login(c *gin.Context) {
-	state := com.RandomString(6)
+	state := utils.RandStringBytesMaskImprSrcUnsafe(6)
 	dao.Cache.Set(fmt.Sprintf("%s%s", model.CtxKeyOauth2State, c.ClientIP()), state, 0)
 	url := oa.oauth2Config.AuthCodeURL(state, oauth2.AccessTypeOnline)
 	c.Redirect(http.StatusFound, url)

cmd/dashboard/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"time"
 
 	"github.com/patrickmn/go-cache"
@@ -10,13 +11,13 @@ import (
 	"github.com/naiba/nezha/cmd/dashboard/controller"
 	"github.com/naiba/nezha/cmd/dashboard/rpc"
 	"github.com/naiba/nezha/model"
+	pb "github.com/naiba/nezha/proto"
 	"github.com/naiba/nezha/service/alertmanager"
 	"github.com/naiba/nezha/service/dao"
 )
 
 func init() {
 	var err error
-	dao.ServerList = make(map[uint64]*model.Server)
 	dao.Conf = &model.Config{}
 	err = dao.Conf.Read("data/config.yaml")
 	if err != nil {
@@ -30,14 +31,26 @@ func init() {
 		dao.DB = dao.DB.Debug()
 	}
 	dao.Cache = cache.New(5*time.Minute, 10*time.Minute)
-	initDB()
+	initSystem()
 }
 
-func initDB() {
+func initSystem() {
 	dao.DB.AutoMigrate(model.Server{}, model.User{},
 		model.Notification{}, model.AlertRule{}, model.Monitor{},
-		model.MonitorHistory{})
-	// load cache
+		model.MonitorHistory{}, model.Cron{})
+
+	loadServers() //加载服务器列表
+	loadCrons()   //加载计划任务
+
+	// 清理旧数据
+	dao.Cron.AddFunc("* 3 * * *", cleanMonitorHistory)
+}
+
+func cleanMonitorHistory() {
+	dao.DB.Delete(&model.MonitorHistory{}, "created_at < ?", time.Now().AddDate(0, -1, 0))
+}
+
+func loadServers() {
 	var servers []model.Server
 	dao.DB.Find(&servers)
 	for _, s := range servers {
@@ -49,6 +62,35 @@ func initDB() {
 	dao.ReSortServer()
 }
 
+func loadCrons() {
+	var crons []model.Cron
+	dao.DB.Find(&crons)
+	var err error
+	for i := 0; i < len(crons); i++ {
+		cr := crons[i]
+		cr.CronID, err = dao.Cron.AddFunc(cr.Scheduler, func() {
+			dao.ServerLock.RLock()
+			defer dao.ServerLock.RUnlock()
+			for j := 0; j < len(cr.Servers); j++ {
+				if dao.ServerList[cr.Servers[j]].TaskStream != nil {
+					dao.ServerList[cr.Servers[j]].TaskStream.Send(&pb.Task{
+						Id:   cr.ID,
+						Data: cr.Command,
+						Type: model.TaskTypeCommand,
+					})
+				} else {
+					alertmanager.SendNotification(fmt.Sprintf("计划任务：%s，服务器：%d 离线，无法执行。", cr.Name, cr.Servers[j]))
+				}
+			}
+		})
+		if err != nil {
+			panic(err)
+		}
+		dao.Crons[cr.ID] = &cr
+	}
+	dao.Cron.Start()
+}
+
 func main() {
 	go controller.ServeWeb(dao.Conf.HTTPPort)
 	go rpc.ServeRPC(5555)