133 lines
3.4 KiB
Go
133 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"net"
|
|
"time"
|
|
_ "time/tzdata"
|
|
|
|
"github.com/ory/graceful"
|
|
"github.com/soheilhy/cmux"
|
|
flag "github.com/spf13/pflag"
|
|
"golang.org/x/crypto/bcrypt"
|
|
|
|
"github.com/naiba/nezha/cmd/dashboard/controller"
|
|
"github.com/naiba/nezha/cmd/dashboard/rpc"
|
|
"github.com/naiba/nezha/model"
|
|
"github.com/naiba/nezha/proto"
|
|
"github.com/naiba/nezha/service/singleton"
|
|
)
|
|
|
|
type DashboardCliParam struct {
|
|
Version bool // 当前版本号
|
|
ConfigFile string // 配置文件路径
|
|
DatebaseLocation string // Sqlite3 数据库文件路径
|
|
}
|
|
|
|
var (
|
|
dashboardCliParam DashboardCliParam
|
|
)
|
|
|
|
func init() {
|
|
flag.CommandLine.ParseErrorsWhitelist.UnknownFlags = true
|
|
flag.BoolVarP(&dashboardCliParam.Version, "version", "v", false, "查看当前版本号")
|
|
flag.StringVarP(&dashboardCliParam.ConfigFile, "config", "c", "data/config.yaml", "配置文件路径")
|
|
flag.StringVar(&dashboardCliParam.DatebaseLocation, "db", "data/sqlite.db", "Sqlite3数据库文件路径")
|
|
flag.Parse()
|
|
|
|
// 初始化 dao 包
|
|
singleton.InitConfigFromPath(dashboardCliParam.ConfigFile)
|
|
singleton.InitTimezoneAndCache()
|
|
singleton.InitDBFromPath(dashboardCliParam.DatebaseLocation)
|
|
initSystem()
|
|
}
|
|
|
|
func initSystem() {
|
|
// 初始化管理员账户
|
|
var usersCount int64
|
|
if err := singleton.DB.Model(&model.User{}).Count(&usersCount).Error; err != nil {
|
|
panic(err)
|
|
}
|
|
if usersCount == 0 {
|
|
hash, err := bcrypt.GenerateFromPassword([]byte("admin"), bcrypt.DefaultCost)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
admin := model.User{
|
|
Username: "admin",
|
|
Password: string(hash),
|
|
}
|
|
if err := singleton.DB.Create(&admin).Error; err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
// 启动 singleton 包下的所有服务
|
|
singleton.LoadSingleton()
|
|
|
|
// 每天的3:30 对 监控记录 和 流量记录 进行清理
|
|
if _, err := singleton.Cron.AddFunc("0 30 3 * * *", singleton.CleanMonitorHistory); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
// 每小时对流量记录进行打点
|
|
if _, err := singleton.Cron.AddFunc("0 0 * * * *", singleton.RecordTransferHourlyUsage); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
func main() {
|
|
if dashboardCliParam.Version {
|
|
fmt.Println(singleton.Version)
|
|
return
|
|
}
|
|
|
|
l, err := net.Listen("tcp", fmt.Sprintf(":%d", singleton.Conf.ListenPort))
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
m := cmux.New(l)
|
|
grpcL := m.Match(cmux.HTTP2HeaderField("content-type", "application/grpc"))
|
|
httpL := m.Match(cmux.HTTP1Fast())
|
|
|
|
singleton.CleanMonitorHistory()
|
|
go rpc.ServeRPC(grpcL)
|
|
serviceSentinelDispatchBus := make(chan model.Monitor) // 用于传递服务监控任务信息的channel
|
|
go rpc.DispatchTask(serviceSentinelDispatchBus)
|
|
go rpc.DispatchKeepalive()
|
|
go singleton.AlertSentinelStart()
|
|
singleton.NewServiceSentinel(serviceSentinelDispatchBus)
|
|
srv := controller.ServeWeb()
|
|
|
|
go dispatchReportInfoTask()
|
|
if err := graceful.Graceful(func() error {
|
|
return srv.Serve(httpL)
|
|
}, func(c context.Context) error {
|
|
log.Println("NEZHA>> Graceful::START")
|
|
singleton.RecordTransferHourlyUsage()
|
|
log.Println("NEZHA>> Graceful::END")
|
|
m.Close()
|
|
return nil
|
|
}); err != nil {
|
|
log.Printf("NEZHA>> ERROR: %v", err)
|
|
}
|
|
}
|
|
|
|
func dispatchReportInfoTask() {
|
|
time.Sleep(time.Second * 15)
|
|
singleton.ServerLock.RLock()
|
|
defer singleton.ServerLock.RUnlock()
|
|
for _, server := range singleton.ServerList {
|
|
if server == nil || server.TaskStream == nil {
|
|
continue
|
|
}
|
|
server.TaskStream.Send(&proto.Task{
|
|
Type: model.TaskTypeReportHostInfo,
|
|
Data: "",
|
|
})
|
|
}
|
|
}
|