From 0a890a20217a1c34181079c26309fa3685cc2353 Mon Sep 17 00:00:00 2001 From: UUBulb <35923940+uubulb@users.noreply.github.com> Date: Wed, 23 Oct 2024 22:34:59 +0800 Subject: [PATCH] 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 | 36 +++++++-- cmd/agent/commands/service.go | 37 +++++++++ cmd/agent/main.go | 128 ++++++++++++++++++++----------- cmd/agent/service.go | 42 ---------- go.mod | 5 ++ go.sum | 10 +++ model/config.go | 49 +++++++----- 7 files changed, 193 insertions(+), 114 deletions(-) rename cmd/agent/{ => commands}/edit.go (74%) create mode 100644 cmd/agent/commands/service.go delete mode 100644 cmd/agent/service.go diff --git a/cmd/agent/edit.go b/cmd/agent/commands/edit.go similarity index 74% rename from cmd/agent/edit.go rename to cmd/agent/commands/edit.go index e2b8d88..3dc6eb9 100644 --- a/cmd/agent/edit.go +++ b/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) diff --git a/cmd/agent/commands/service.go b/cmd/agent/commands/service.go new file mode 100644 index 0000000..9ff1b51 --- /dev/null +++ b/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() +} diff --git a/cmd/agent/main.go b/cmd/agent/main.go index 13269e0..b9e925d 100644 --- a/cmd/agent/main.go +++ b/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", "", "服务操作 ") - 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: "", + 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) diff --git a/cmd/agent/service.go b/cmd/agent/service.go deleted file mode 100644 index 0fe66d4..0000000 --- a/cmd/agent/service.go +++ /dev/null @@ -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() -} diff --git a/go.mod b/go.mod index 9b07ade..cdd5cfc 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/dean2021/goss v0.0.0-20230129073947-df90431348f1 github.com/ebi-yade/altsvc-go v0.1.1 github.com/ebitengine/purego v0.8.0 + github.com/hashicorp/go-uuid v1.0.3 github.com/iamacarpet/go-winpty v1.0.4 github.com/jaypipes/ghw v0.12.0 github.com/json-iterator/go v1.1.12 @@ -21,6 +22,7 @@ require ( github.com/refraction-networking/utls v1.6.3 github.com/shirou/gopsutil/v4 v4.24.9 github.com/spf13/viper v1.19.0 + github.com/urfave/cli/v2 v2.27.5 golang.org/x/net v0.29.0 golang.org/x/sys v0.25.0 google.golang.org/grpc v1.64.1 @@ -34,6 +36,7 @@ require ( github.com/andybalholm/brotli v1.0.6 // indirect github.com/antihax/optional v1.0.0 // indirect github.com/cloudflare/circl v1.3.7 // indirect + github.com/cpuguy83/go-md2man/v2 v2.0.5 // indirect github.com/fsnotify/fsnotify v1.7.0 // indirect github.com/ghodss/yaml v1.0.0 // indirect github.com/go-ole/go-ole v1.2.6 // indirect @@ -62,6 +65,7 @@ require ( github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b // indirect github.com/quic-go/qpack v0.4.0 // indirect github.com/quic-go/qtls-go1-20 v0.4.1 // indirect + github.com/russross/blackfriday/v2 v2.1.0 // indirect github.com/sagikazarmark/locafero v0.4.0 // indirect github.com/sagikazarmark/slog-shim v0.1.0 // indirect github.com/sourcegraph/conc v0.3.0 // indirect @@ -73,6 +77,7 @@ require ( github.com/tklauser/go-sysconf v0.3.12 // indirect github.com/tklauser/numcpus v0.6.1 // indirect github.com/ulikunitz/xz v0.5.11 // indirect + github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1 // indirect github.com/yusufpapurcu/wmi v1.2.4 // indirect go.uber.org/atomic v1.9.0 // indirect go.uber.org/mock v0.4.0 // indirect diff --git a/go.sum b/go.sum index 9f2c63c..51ebbe0 100644 --- a/go.sum +++ b/go.sum @@ -22,6 +22,8 @@ github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5P github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= github.com/cloudflare/circl v1.3.7 h1:qlCDlTPz2n9fu58M0Nh1J/JzcFpfgkFHHX3O35r5vcU= github.com/cloudflare/circl v1.3.7/go.mod h1:sRTcRWXGLrKw6yIGJ+l7amYJFfAXbZG0kBSc8r4zxgA= +github.com/cpuguy83/go-md2man/v2 v2.0.5 h1:ZtcqGrnekaHpVLArFSe4HK5DoKx1T0rq2DwVB0alcyc= +github.com/cpuguy83/go-md2man/v2 v2.0.5/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= github.com/creack/pty v1.1.17/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4= github.com/creack/pty v1.1.23 h1:4M6+isWdcStXEf15G/RbrMPOQj1dZ7HPZCGwE4kOeP0= github.com/creack/pty v1.1.23/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= @@ -59,6 +61,8 @@ github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 h1:yAJXTCF9TqKcTiHJAE github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/go-uuid v1.0.3 h1:2gKiV6YVmrJ1i2CKKa9obLvRieoRGviZFL26PcT/Co8= +github.com/hashicorp/go-uuid v1.0.3/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec h1:qv2VnGeEQHchGaZ/u7lxST/RaJw+cv273q79D81Xbog= @@ -130,6 +134,8 @@ github.com/quic-go/quic-go v0.40.1/go.mod h1:PeN7kuVJ4xZbxSv/4OX6S1USOX8MJvydwpT github.com/refraction-networking/utls v1.6.3 h1:MFOfRN35sSx6K5AZNIoESsBuBxS2LCgRilRIdHb6fDc= github.com/refraction-networking/utls v1.6.3/go.mod h1:yil9+7qSl+gBwJqztoQseO6Pr3h62pQoY1lXiNR/FPs= github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= +github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6keLGt6kNQ= github.com/sagikazarmark/locafero v0.4.0/go.mod h1:Pe1W6UlPYUk/+wc/6KFhbORCfqzgYEpgQ3O5fPuL3H4= github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE= @@ -167,6 +173,10 @@ github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+F github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY= github.com/ulikunitz/xz v0.5.11 h1:kpFauv27b6ynzBNT/Xy+1k+fK4WswhN/6PN5WhFAGw8= github.com/ulikunitz/xz v0.5.11/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= +github.com/urfave/cli/v2 v2.27.5 h1:WoHEJLdsXr6dDWoJgMq/CboDmyY/8HMMH1fTECbih+w= +github.com/urfave/cli/v2 v2.27.5/go.mod h1:3Sevf16NykTbInEnD0yKkjDAeZDS0A6bzhBH5hrMvTQ= +github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1 h1:gEOO8jv9F4OT7lGCjxCBTO/36wtF6j2nSip77qHd4x4= +github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1/go.mod h1:Ohn+xnUBiLI6FVj/9LpzZWtj1/D6lUovWYBkxHVV3aM= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0= github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= diff --git a/model/config.go b/model/config.go index 8969644..8060fcb 100644 --- a/model/config.go +++ b/model/config.go @@ -8,28 +8,28 @@ import ( ) type AgentConfig struct { - Debug bool `mapstructure:"debug"` + Debug bool `mapstructure:"debug" json:"debug"` - Server string `mapstructure:"server"` // 服务器地址 - ClientSecret string `mapstructure:"client_secret"` // 客户端密钥 - UUID string `mapstructure:"uuid"` + Server string `mapstructure:"server" json:"server"` // 服务器地址 + ClientSecret string `mapstructure:"client_secret" json:"client_secret"` // 客户端密钥 + UUID string `mapstructure:"uuid" json:"uuid"` - HardDrivePartitionAllowlist []string `mapstructure:"hard_drive_partition_allowlist"` - NICAllowlist map[string]bool `mapstructure:"nic_allowlist"` - DNS []string `mapstructure:"dns"` - GPU bool `mapstructure:"gpu"` // 是否检查GPU - Temperature bool `mapstructure:"temperature"` // 是否检查温度 - SkipConnectionCount bool `mapstructure:"skip_connection_count"` // 跳过连接数检查 - SkipProcsCount bool `mapstructure:"skip_procs_count"` // 跳过进程数量检查 - DisableAutoUpdate bool `mapstructure:"disable_auto_update"` // 关闭自动更新 - DisableForceUpdate bool `mapstructure:"disable_force_update"` // 关闭强制更新 - DisableCommandExecute bool `mapstructure:"disable_command_execute"` // 关闭命令执行 - ReportDelay int `mapstructure:"report_delay"` // 报告间隔 - TLS bool `mapstructure:"tls"` // 是否使用TLS加密传输至服务端 - InsecureTLS bool `mapstructure:"insecure_tls"` // 是否禁用证书检查 - UseIPv6CountryCode bool `mapstructure:"use_i_pv_6_country_code"` // 默认优先展示IPv6旗帜 - UseGiteeToUpgrade bool `mapstructure:"use_gitee_to_upgrade"` // 强制从Gitee获取更新 - IPReportPeriod uint32 `mapstructure:"ip_report_period"` // IP上报周期 + HardDrivePartitionAllowlist []string `mapstructure:"hard_drive_partition_allowlist" json:"hard_drive_partition_allowlist"` + NICAllowlist map[string]bool `mapstructure:"nic_allowlist" json:"nic_allowlist"` + DNS []string `mapstructure:"dns" json:"dns"` + GPU bool `mapstructure:"gpu" json:"gpu"` // 是否检查GPU + Temperature bool `mapstructure:"temperature" json:"temperature"` // 是否检查温度 + SkipConnectionCount bool `mapstructure:"skip_connection_count" json:"skip_connection_count"` // 跳过连接数检查 + SkipProcsCount bool `mapstructure:"skip_procs_count" json:"skip_procs_count"` // 跳过进程数量检查 + DisableAutoUpdate bool `mapstructure:"disable_auto_update" json:"disable_auto_update"` // 关闭自动更新 + DisableForceUpdate bool `mapstructure:"disable_force_update" json:"disable_force_update"` // 关闭强制更新 + DisableCommandExecute bool `mapstructure:"disable_command_execute" json:"disable_command_execute"` // 关闭命令执行 + ReportDelay int `mapstructure:"report_delay" json:"report_delay"` // 报告间隔 + TLS bool `mapstructure:"tls" json:"tls"` // 是否使用TLS加密传输至服务端 + InsecureTLS bool `mapstructure:"insecure_tls" json:"insecure_tls"` // 是否禁用证书检查 + UseIPv6CountryCode bool `mapstructure:"use_ipv6_country_code" json:"use_ipv6_country_code"` // 默认优先展示IPv6旗帜 + UseGiteeToUpgrade bool `mapstructure:"use_gitee_to_upgrade" json:"use_gitee_to_upgrade"` // 强制从Gitee获取更新 + IPReportPeriod uint32 `mapstructure:"ip_report_period" json:"ip_report_period"` // IP上报周期 v *viper.Viper } @@ -46,6 +46,15 @@ func (c *AgentConfig) Read(path string) error { if err != nil { return err } + + if c.ReportDelay == 0 { + c.ReportDelay = 1 + } + + if c.IPReportPeriod == 0 { + c.IPReportPeriod = 1800 + } + return nil }