Skip to content

Commit

Permalink
Update.
Browse files Browse the repository at this point in the history
  • Loading branch information
@xOS
xOS committed Mar 19, 2022
1 parent de78957 commit 1fce71c
Show file tree
Hide file tree
Showing 5 changed files with 233 additions and 92 deletions.
148 changes: 116 additions & 32 deletions cmd/agent/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,22 @@ import (
"io"
"net/http"
"os"
"path/filepath"
"runtime"
"strings"
"time"

"github.com/AlecAivazis/survey/v2"
"github.com/blang/semver"
"github.com/gorilla/websocket"
"github.com/p14yground/go-github-selfupdate/selfupdate"
"github.com/shirou/gopsutil/v3/disk"
"github.com/shirou/gopsutil/v3/host"
psnet "github.com/shirou/gopsutil/v3/net"
flag "github.com/spf13/pflag"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure"

"github.com/xos/serverstatus/cmd/agent/monitor"
"github.com/xos/serverstatus/cmd/agent/pty"
Expand All @@ -26,7 +32,7 @@ import (
"github.com/xos/serverstatus/service/rpc"
)

type AgentConfig struct {
type AgentCliParam struct {
SkipConnectionCount bool
SkipProcsCount bool
DisableAutoUpdate bool
Expand All @@ -47,9 +53,10 @@ var (
)

var (
agentConf AgentConfig
updateCh = make(chan struct{}) // Agent 自动更新间隔
httpClient = &http.Client{
agentCliParam AgentCliParam
agentConfig model.AgentConfig
updateCh = make(chan struct{}) // Agent 自动更新间隔
httpClient = &http.Client{
CheckRedirect: func(req *http.Request, via []*http.Request) error {
return http.ErrUseLastResponse
},
Expand All @@ -65,6 +72,12 @@ const (
func init() {
http.DefaultClient.Timeout = time.Second * 5
flag.CommandLine.ParseErrorsWhitelist.UnknownFlags = true

ex, err := os.Executable()
if err != nil {
panic(err)
}
agentConfig.Read(filepath.Dir(ex) + "/config.yml")
}

func main() {
Expand All @@ -90,24 +103,31 @@ func main() {
// 来自于 GoReleaser 的版本号
monitor.Version = version

flag.BoolVarP(&agentConf.Debug, "debug", "d", true, "开启调试信息")
flag.StringVarP(&agentConf.Server, "server", "s", "localhost:2222", "管理面板GRPC端口")
flag.StringVarP(&agentConf.ClientSecret, "password", "p", "", "探针连接密钥")
flag.IntVar(&agentConf.ReportDelay, "report-delay", 1, "系统状态上报间隔")
flag.BoolVar(&agentConf.SkipConnectionCount, "skip-conn", false, "不监控连接数")
flag.BoolVar(&agentConf.SkipProcsCount, "skip-procs", false, "不监控进程数")
flag.BoolVar(&agentConf.DisableCommandExecute, "disable-command-execute", false, "禁止在此机器上执行命令")
flag.BoolVar(&agentConf.DisableAutoUpdate, "disable-auto-update", false, "禁用自动升级")
flag.BoolVar(&agentConf.DisableForceUpdate, "disable-force-update", false, "禁用强制升级")
flag.BoolVar(&agentConf.TLS, "tls", false, "启用SSL/TLS加密")
var isEditAgentConfig bool
flag.BoolVarP(&agentCliParam.Debug, "debug", "d", true, "开启调试信息")
flag.BoolVarP(&isEditAgentConfig, "edit-agent-config", "", false, "修改要监控的网卡/分区白名单")
flag.StringVarP(&agentCliParam.Server, "server", "s", "localhost:2222", "管理面板RPC端口")
flag.StringVarP(&agentCliParam.ClientSecret, "password", "p", "", "探针连接Secret")
flag.IntVar(&agentCliParam.ReportDelay, "report-delay", 1, "系统状态上报间隔")
flag.BoolVar(&agentCliParam.SkipConnectionCount, "skip-conn", false, "不监控连接数")
flag.BoolVar(&agentCliParam.SkipProcsCount, "skip-procs", false, "不监控进程数")
flag.BoolVar(&agentCliParam.DisableCommandExecute, "disable-command-execute", false, "禁止在此机器上执行命令")
flag.BoolVar(&agentCliParam.DisableAutoUpdate, "disable-auto-update", false, "禁用自动升级")
flag.BoolVar(&agentCliParam.DisableForceUpdate, "disable-force-update", false, "禁用强制升级")
flag.BoolVar(&agentCliParam.TLS, "tls", false, "启用SSL/TLS加密")
flag.Parse()

if agentConf.ClientSecret == "" {
if isEditAgentConfig {
editAgentConfig()
return
}

if agentCliParam.ClientSecret == "" {
flag.Usage()
return
}

if agentConf.ReportDelay < 1 || agentConf.ReportDelay > 4 {
if agentCliParam.ReportDelay < 1 || agentCliParam.ReportDelay > 4 {
println("report-delay 的区间为 1-4")
return
}
Expand All @@ -117,18 +137,18 @@ func main() {

func run() {
auth := rpc.AuthHandler{
ClientSecret: agentConf.ClientSecret,
ClientSecret: agentCliParam.ClientSecret,
}

if !agentConf.DisableCommandExecute {
if !agentCliParam.DisableCommandExecute {
go pty.DownloadDependency()
}
// 上报服务器信息
go reportState()
// 更新IP信息
go monitor.UpdateIP()

if _, err := semver.Parse(version); err == nil && !agentConf.DisableAutoUpdate {
if _, err := semver.Parse(version); err == nil && !agentCliParam.DisableAutoUpdate {
go func() {
for range updateCh {
go func() {
Expand Down Expand Up @@ -159,12 +179,12 @@ func run() {
for {
timeOutCtx, cancel := context.WithTimeout(context.Background(), networkTimeOut)
var securityOption grpc.DialOption
if agentConf.TLS {
if agentCliParam.TLS {
securityOption = grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{MinVersion: tls.VersionTLS12}))
} else {
securityOption = grpc.WithInsecure()
securityOption = grpc.WithTransportCredentials(insecure.NewCredentials())
}
conn, err = grpc.DialContext(timeOutCtx, agentConf.Server, securityOption, grpc.WithPerRPCCredentials(&auth))
conn, err = grpc.DialContext(timeOutCtx, agentCliParam.Server, securityOption, grpc.WithPerRPCCredentials(&auth))
if err != nil {
println("与面板建立连接失败：", err)
cancel()
Expand All @@ -175,7 +195,7 @@ func run() {
client = pb.NewProbeServiceClient(conn)
// 第一步注册
timeOutCtx, cancel = context.WithTimeout(context.Background(), networkTimeOut)
_, err = client.ReportSystemInfo(timeOutCtx, monitor.GetHost().PB())
_, err = client.ReportSystemInfo(timeOutCtx, monitor.GetHost(&agentConfig).PB())
if err != nil {
println("上报系统信息失败：", err)
cancel()
Expand All @@ -185,7 +205,7 @@ func run() {
cancel()
inited = true
// 执行 Task
tasks, err := client.RequestTask(context.Background(), monitor.GetHost().PB())
tasks, err := client.RequestTask(context.Background(), monitor.GetHost(&agentConfig).PB())
if err != nil {
println("请求任务失败：", err)
retry()
Expand Down Expand Up @@ -240,20 +260,20 @@ func reportState() {
for {
// 为了更准确的记录时段流量，inited 后再上传状态信息
if client != nil && inited {
monitor.TrackNetworkSpeed()
monitor.TrackNetworkSpeed(&agentConfig)
timeOutCtx, cancel := context.WithTimeout(context.Background(), networkTimeOut)
_, err = client.ReportSystemState(timeOutCtx, monitor.GetState(agentConf.SkipConnectionCount, agentConf.SkipProcsCount).PB())
_, err = client.ReportSystemState(timeOutCtx, monitor.GetState(&agentConfig, agentCliParam.SkipConnectionCount, agentCliParam.SkipProcsCount).PB())
cancel()
if err != nil {
println("reportState error", err)
time.Sleep(delayWhenError)
}
if lastReportHostInfo.Before(time.Now().Add(-10 * time.Minute)) {
lastReportHostInfo = time.Now()
client.ReportSystemInfo(context.Background(), monitor.GetHost().PB())
client.ReportSystemInfo(context.Background(), monitor.GetHost(&agentConfig).PB())
}
}
time.Sleep(time.Second * time.Duration(agentConf.ReportDelay))
time.Sleep(time.Second * time.Duration(agentCliParam.ReportDelay))
}
}

Expand All @@ -274,7 +294,7 @@ func doSelfUpdate(useLocalVersion bool) {
}

func handleUpgradeTask(task *pb.Task, result *pb.TaskResult) {
if agentConf.DisableForceUpdate {
if agentCliParam.DisableForceUpdate {
return
}
doSelfUpdate(false)
Expand All @@ -286,7 +306,7 @@ type WindowSize struct {
}

func handleTerminalTask(task *pb.Task) {
if agentConf.DisableCommandExecute {
if agentCliParam.DisableCommandExecute {
println("该节点已禁止命令执行")
return
}
Expand All @@ -301,7 +321,7 @@ func handleTerminalTask(task *pb.Task) {
protocol += "s"
}
header := http.Header{}
header.Add("Secret", agentConf.ClientSecret)
header.Add("Secret", agentCliParam.ClientSecret)
conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s://%s/terminal/%s", protocol, terminal.Host, terminal.Session), header)
if err != nil {
println("Terminal 连接失败：", err)
Expand Down Expand Up @@ -370,8 +390,72 @@ func handleTerminalTask(task *pb.Task) {
}
}

func editAgentConfig() {
nc, err := psnet.IOCounters(true)
if err != nil {
panic(err)
}
var nicAllowlistOptions []string
for _, v := range nc {
nicAllowlistOptions = append(nicAllowlistOptions, v.Name)
}

var diskAllowlistOptions []string
diskList, err := disk.Partitions(false)
if err != nil {
panic(err)
}
for _, p := range diskList {
diskAllowlistOptions = append(diskAllowlistOptions, fmt.Sprintf("%s\t%s\t%s", p.Mountpoint, p.Fstype, p.Device))
}

var qs = []*survey.Question{
{
Name: "nic",
Prompt: &survey.MultiSelect{
Message: "选择要监控的网卡",
Options: nicAllowlistOptions,
},
},
{
Name: "disk",
Prompt: &survey.MultiSelect{
Message: "选择要监控的硬盘分区",
Options: diskAllowlistOptions,
},
},
}

answers := struct {
Nic []string
Disk []string
}{}

err = survey.Ask(qs, &answers, survey.WithValidator(survey.Required))
if err != nil {
fmt.Println("选择错误", err.Error())
return
}

agentConfig.HardDrivePartitionAllowlist = []string{}
for _, v := range answers.Disk {
agentConfig.HardDrivePartitionAllowlist = append(agentConfig.HardDrivePartitionAllowlist, strings.Split(v, "\t")[0])
}

agentConfig.NICAllowlist = make(map[string]bool)
for _, v := range answers.Nic {
agentConfig.NICAllowlist[v] = true
}

if err = agentConfig.Save(); err != nil {
panic(err)
}

fmt.Println("修改自定义配置成功，重启 Agnet 后生效")
}

func println(v ...interface{}) {
if agentConf.Debug {
if agentCliParam.Debug {
fmt.Printf("NG@%s>> ", time.Now().Format("2006-01-02 15:04:05"))
fmt.Println(v...)
}
Expand Down
Loading

0 comments on commit 1fce71c

Please sign in to comment.