make edit and service flag subcommands (#77)

* make edit and service argument subcommands

* generate uuid if non-exist, use default report_delay or ip_report_period value if not specified

cmd/agent/commands/edit.go

@@ -1,4 +1,4 @@
-package main
+package commands
 
 import (
 	"errors"
@@ -7,12 +7,15 @@ import (
 	"strings"
 
 	"github.com/AlecAivazis/survey/v2"
+	"github.com/hashicorp/go-uuid"
 	"github.com/shirou/gopsutil/v4/disk"
 	psnet "github.com/shirou/gopsutil/v4/net"
+
+	"github.com/nezhahq/agent/model"
 )
 
 // 修改Agent要监控的网卡与硬盘分区
-func editAgentConfig(configPath string) {
+func EditAgentConfig(configPath string, agentConfig *model.AgentConfig) {
 	agentConfig.Read(configPath)
 
 	nc, err := psnet.IOCounters(true)
@@ -33,6 +36,11 @@ func editAgentConfig(configPath string) {
 		diskAllowlistOptions = append(diskAllowlistOptions, fmt.Sprintf("%s\t%s\t%s", p.Mountpoint, p.Fstype, p.Device))
 	}
 
+	uuid, err := uuid.GenerateUUID()
+	if err != nil {
+		panic(err)
+	}
+
 	var qs = []*survey.Question{
 		{
 			Name: "nic",
@@ -55,6 +63,16 @@ func editAgentConfig(configPath string) {
 				Default: strings.Join(agentConfig.DNS, ","),
 			},
 		},
+		{
+			Name: "uuid",
+			Prompt: &survey.Input{
+				Message: "输入 Agent UUID",
+				Default: agentConfig.UUID,
+				Suggest: func(_ string) []string {
+					return []string{uuid}
+				},
+			},
+		},
 		{
 			Name: "gpu",
 			Prompt: &survey.Confirm{
@@ -79,12 +97,13 @@ func editAgentConfig(configPath string) {
 	}
 
 	answers := struct {
-		Nic         []string
-		Disk        []string
-		DNS         string
-		GPU         bool
-		Temperature bool
-		Debug       bool
+		Nic         []string `mapstructure:"nic_allowlist" json:"nic_allowlist"`
+		Disk        []string `mapstructure:"hard_drive_partition_allowlist" json:"hard_drive_partition_allowlist"`
+		DNS         string   `mapstructure:"dns" json:"dns"`
+		GPU         bool     `mapstructure:"gpu" json:"gpu"`
+		Temperature bool     `mapstructure:"temperature" json:"temperature"`
+		Debug       bool     `mapstructure:"debug" json:"debug"`
+		UUID        string   `mapstructure:"uuid" json:"uuid"`
 	}{}
 
 	err = survey.Ask(qs, &answers, survey.WithValidator(survey.Required))
@@ -125,6 +144,7 @@ func editAgentConfig(configPath string) {
 	agentConfig.GPU = answers.GPU
 	agentConfig.Temperature = answers.Temperature
 	agentConfig.Debug = answers.Debug
+	agentConfig.UUID = answers.UUID
 
 	if err = agentConfig.Save(); err != nil {
 		panic(err)

cmd/agent/commands/service.go

@@ -0,0 +1,37 @@
+package commands
+
+import (
+	"os"
+
+	"github.com/nezhahq/service"
+)
+
+type Program struct {
+	Exit    chan struct{}
+	Service service.Service
+	Run     func()
+}
+
+func (p *Program) Start(s service.Service) error {
+	go p.run()
+	return nil
+}
+
+func (p *Program) Stop(s service.Service) error {
+	close(p.Exit)
+	if service.Interactive() {
+		os.Exit(0)
+	}
+	return nil
+}
+
+func (p *Program) run() {
+	defer func() {
+		if service.Interactive() {
+			p.Stop(p.Service)
+		} else {
+			p.Service.Stop()
+		}
+	}()
+	p.Run()
+}

cmd/agent/main.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"crypto/tls"
 	"errors"
-	"flag"
 	"fmt"
 	"io"
 	"log"
@@ -20,17 +19,20 @@ import (
 
 	"github.com/blang/semver"
 	"github.com/ebi-yade/altsvc-go"
+	"github.com/hashicorp/go-uuid"
 	"github.com/nezhahq/go-github-selfupdate/selfupdate"
 	"github.com/nezhahq/service"
 	ping "github.com/prometheus-community/pro-bing"
 	"github.com/quic-go/quic-go/http3"
 	utls "github.com/refraction-networking/utls"
 	"github.com/shirou/gopsutil/v4/host"
+	"github.com/urfave/cli/v2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/resolver"
 
+	"github.com/nezhahq/agent/cmd/agent/commands"
 	"github.com/nezhahq/agent/model"
 	fm "github.com/nezhahq/agent/pkg/fm"
 	"github.com/nezhahq/agent/pkg/monitor"
@@ -42,14 +44,15 @@ import (
 )
 
 var (
-	version        string
-	arch           string
-	executablePath string
-	client         pb.NezhaServiceClient
-	initialized    bool
-	dnsResolver    = &net.Resolver{PreferGo: true}
-	agentConfig    model.AgentConfig
-	httpClient     = &http.Client{
+	version           string
+	arch              string
+	defaultConfigPath string
+	executablePath    string
+	client            pb.NezhaServiceClient
+	initialized       bool
+	dnsResolver       = &net.Resolver{PreferGo: true}
+	agentConfig       model.AgentConfig
+	httpClient        = &http.Client{
 		CheckRedirect: func(req *http.Request, via []*http.Request) error {
 			return http.ErrUseLastResponse
 		},
@@ -101,42 +104,74 @@ func init() {
 
 	// 来自于 GoReleaser 的版本号
 	monitor.Version = version
-}
 
-func main() {
 	var err error
 	executablePath, err = os.Executable()
 	if err != nil {
 		panic(err)
 	}
 
-	var showVersion, isEditConfig, showHelp bool
-	var configPath, serviceAction string
+	defaultConfigPath = filepath.Join(filepath.Dir(executablePath), "config.yml")
+}
 
-	// 初始化运行参数
-	flag.BoolVar(&showVersion, "v", false, "查看当前版本号")
-	flag.BoolVar(&showHelp, "h", false, "查看帮助")
-	flag.BoolVar(&isEditConfig, "edit", false, "编辑配置文件")
-	flag.StringVar(&serviceAction, "service", "", "服务操作 <install/uninstall/start/stop/restart>")
-	flag.StringVar(&configPath, "c", filepath.Dir(executablePath)+"/config.yml", "配置文件路径")
-
-	flag.Parse()
-
-	if showHelp {
-		flag.Usage()
-		os.Exit(0)
+func main() {
+	app := &cli.App{
+		Usage:   "哪吒监控 Agent",
+		Version: version,
+		Flags: []cli.Flag{
+			&cli.BoolFlag{Name: "version", Aliases: []string{"v"}, Usage: "查看当前版本号"},
+			&cli.StringFlag{Name: "config", Aliases: []string{"c"}, Usage: "配置文件路径"},
+		},
+		Action: func(c *cli.Context) error {
+			if c.Bool("version") {
+				fmt.Println(c.App.Version)
+				return nil
+			}
+			if path := c.String("config"); path != "" {
+				preRun(path)
+			} else {
+				preRun(defaultConfigPath)
+			}
+			runService("")
+			return nil
+		},
+		Commands: []*cli.Command{
+			{
+				Name:  "edit",
+				Usage: "编辑配置文件",
+				Flags: []cli.Flag{
+					&cli.StringFlag{Name: "config", Aliases: []string{"c"}, Usage: "配置文件路径"},
+				},
+				Action: func(c *cli.Context) error {
+					if path := c.String("config"); path != "" {
+						commands.EditAgentConfig(path, &agentConfig)
+					} else {
+						commands.EditAgentConfig(defaultConfigPath, &agentConfig)
+					}
+					return nil
+				},
+			},
+			{
+				Name:      "service",
+				Usage:     "服务操作",
+				UsageText: "<install/uninstall/start/stop/restart>",
+				Action: func(c *cli.Context) error {
+					if arg := c.Args().Get(0); arg != "" {
+						runService(arg)
+						return nil
+					}
+					return cli.Exit("必须指定一个参数", 1)
+				},
+			},
+		},
 	}
 
-	if showVersion {
-		fmt.Println(version)
-		os.Exit(0)
-	}
-
-	if isEditConfig {
-		editAgentConfig(configPath)
-		os.Exit(0)
+	if err := app.Run(os.Args); err != nil {
+		log.Fatal(err)
 	}
+}
 
+func preRun(configPath string) {
 	// windows环境处理
 	if runtime.GOOS == "windows" {
 		hostArch, err := host.KernelArch()
@@ -158,23 +193,27 @@ func main() {
 	}
 
 	if err := agentConfig.Read(configPath); err != nil {
-		println(err)
-		os.Exit(1)
+		log.Fatalf("打开配置文件失败：%v", err)
 	}
 
 	monitor.InitConfig(&agentConfig)
 
 	if agentConfig.ClientSecret == "" {
-		println("ClientSecret 不能为空")
-		os.Exit(1)
+		log.Fatal("ClientSecret 不能为空")
 	}
 
 	if agentConfig.ReportDelay < 1 || agentConfig.ReportDelay > 4 {
-		println("report-delay 的区间为 1-4")
-		os.Exit(1)
+		log.Fatal("report-delay 的区间为 1-4")
 	}
 
-	runService(serviceAction)
+	if agentConfig.UUID == "" {
+		if uuid, err := uuid.GenerateUUID(); err == nil {
+			agentConfig.UUID = uuid
+			agentConfig.Save()
+		} else {
+			log.Fatalf("生成 UUID 失败：%v", err)
+		}
+	}
 }
 
 func run() {
@@ -269,13 +308,14 @@ func runService(action string) {
 	svcConfig := &service.Config{
 		Name:             filepath.Base(executablePath),
 		DisplayName:      filepath.Base(executablePath),
-		Description:      "哪吒探针监控端",
+		Description:      "哪吒监控 Agent",
 		WorkingDirectory: filepath.Dir(executablePath),
 		Option:           winConfig,
 	}
 
-	prg := &program{
-		exit: make(chan struct{}),
+	prg := &commands.Program{
+		Exit: make(chan struct{}),
+		Run:  run,
 	}
 	s, err := service.New(prg, svcConfig)
 	if err != nil {
@@ -283,7 +323,7 @@ func runService(action string) {
 		run()
 		return
 	}
-	prg.service = s
+	prg.Service = s
 
 	if agentConfig.Debug {
 		serviceLogger, err := s.Logger(nil)

cmd/agent/service.go

@@ -1,42 +0,0 @@
-package main
-
-import (
-	"os"
-
-	"github.com/nezhahq/service"
-)
-
-type AgentCliFlags struct {
-	IsSpecified bool
-	Flag        string
-	Value       string
-}
-
-type program struct {
-	exit    chan struct{}
-	service service.Service
-}
-
-func (p *program) Start(s service.Service) error {
-	go p.run()
-	return nil
-}
-
-func (p *program) Stop(s service.Service) error {
-	close(p.exit)
-	if service.Interactive() {
-		os.Exit(0)
-	}
-	return nil
-}
-
-func (p *program) run() {
-	defer func() {
-		if service.Interactive() {
-			p.Stop(p.service)
-		} else {
-			p.service.Stop()
-		}
-	}()
-	run()
-}