KJZH001/nezhahq-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: code style (#44)

Browse Source 
* some code refactor

* service

* chore: code style

---------

Co-authored-by: naiba <hi@nai.ba>
This commit is contained in:
UUBulb 2024-07-23 23:20:55 +08:00 committed by GitHub
parent a99bf7066b
commit c0288971d8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 57 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
cmd/agent/main.go
View File

@ -38,6 +38,7 @@ import (
pb "github.com/nezhahq/agent/proto"
)
// Agent 运行时参数。如需添加新参数，记得同时在 service.go 中添加
type AgentCliParam struct {
SkipConnectionCount bool // 跳过连接数检查
SkipProcsCount bool // 跳过进程数量检查
@ -432,7 +433,9 @@ func reportState() {
// doSelfUpdate 执行更新检查 如果更新成功则会结束进程
func doSelfUpdate(useLocalVersion bool) {
<-monitor.Sync
if monitor.CachedCountry == "" {
return
}
v := semver.MustParse("0.1.0")
if useLocalVersion {
v = semver.MustParse(version)

1
cmd/agent/service.go
View File

@ -86,6 +86,7 @@ func serviceActions(cmd *cobra.Command, args []string) {
{agentCliParam.DisableForceUpdate, "--disable-force-update", ""},
{agentCliParam.UseIPv6CountryCode, "--use-ipv6-countrycode", ""},
{agentConfig.GPU, "--gpu", ""},
{agentCliParam.UseGiteeToUpgrade, "--gitee", ""},
{agentCliParam.IPReportPeriod != 30*60, "-u", fmt.Sprint(agentCliParam.IPReportPeriod)},
}

79
pkg/monitor/monitor.go
View File

@ -83,17 +83,14 @@ func GetHost() *model.Host {
} else {
if hi.VirtualizationRole == "guest" {
cpuType = "Virtual"
ret.Virtualization = hi.VirtualizationSystem
} else {
cpuType = "Physical"
ret.Virtualization = ""
}
ret.Platform = hi.Platform
ret.PlatformVersion = hi.PlatformVersion
ret.Arch = hi.KernelArch
if cpuType == "Physical" {
ret.Virtualization = ""
} else {
ret.Virtualization = hi.VirtualizationSystem
}
ret.BootTime = hi.BootTime
}
@ -218,40 +215,6 @@ func GetState(skipConnectionCount bool, skipProcsCount bool) *model.HostState {
}
}
var tcpConnCount, udpConnCount uint64
if !skipConnectionCount {
ss_err := true
if runtime.GOOS == "linux" {
tcpStat, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_TCP)
udpStat, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_UDP)
if err_tcp == nil && err_udp == nil {
ss_err = false
tcpConnCount = uint64(len(tcpStat))
udpConnCount = uint64(len(udpStat))
}
if strings.Contains(CachedIP, ":") {
tcpStat6, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_TCP)
udpStat6, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_UDP)
if err_tcp == nil && err_udp == nil {
ss_err = false
tcpConnCount += uint64(len(tcpStat6))
udpConnCount += uint64(len(udpStat6))
}
}
}
if ss_err {
conns, _ := net.Connections("all")
for i := 0; i < len(conns); i++ {
switch conns[i].Type {
case syscall.SOCK_STREAM:
tcpConnCount++
case syscall.SOCK_DGRAM:
udpConnCount++
}
}
}
}
if agentConfig.Temperature {
go updateTemperatureStat()
ret.Temperatures = temperatureStat
@ -265,7 +228,7 @@ func GetState(skipConnectionCount bool, skipProcsCount bool) *model.HostState {
ret.NetInTransfer, ret.NetOutTransfer = netInTransfer, netOutTransfer
ret.NetInSpeed, ret.NetOutSpeed = netInSpeed, netOutSpeed
ret.Uptime = uint64(time.Since(cachedBootTime).Seconds())
ret.TcpConnCount, ret.UdpConnCount = tcpConnCount, udpConnCount
ret.TcpConnCount, ret.UdpConnCount = getConns(skipConnectionCount)
return &ret
}
@ -352,6 +315,42 @@ func getDiskTotalAndUsed() (total uint64, used uint64) {
return
}
func getConns(skipConnectionCount bool) (tcpConnCount, udpConnCount uint64) {
if !skipConnectionCount {
ss_err := true
if runtime.GOOS == "linux" {
tcpStat, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_TCP)
udpStat, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_UDP)
if err_tcp == nil && err_udp == nil {
ss_err = false
tcpConnCount = uint64(len(tcpStat))
udpConnCount = uint64(len(udpStat))
}
if strings.Contains(CachedIP, ":") {
tcpStat6, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_TCP)
udpStat6, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_UDP)
if err_tcp == nil && err_udp == nil {
ss_err = false
tcpConnCount += uint64(len(tcpStat6))
udpConnCount += uint64(len(udpStat6))
}
}
}
if ss_err {
conns, _ := net.Connections("all")
for i := 0; i < len(conns); i++ {
switch conns[i].Type {
case syscall.SOCK_STREAM:
tcpConnCount++
case syscall.SOCK_DGRAM:
udpConnCount++
}
}
}
}
return tcpConnCount, udpConnCount
}
func updateGPUStat(gpuStat *uint64) {
if !atomic.CompareAndSwapInt32(&updateGPUStatus, 0, 1) {
return

31
pkg/monitor/myip.go
View File

@ -49,7 +49,6 @@ var (
// "https://freegeoip.app/json/", // 需要 Key
}
CachedIP, CachedCountry string
Sync = make(chan bool)
httpClientV4 = util.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, false)
httpClientV6 = util.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, true)
)
@ -58,16 +57,10 @@ var (
func UpdateIP(useIPv6CountryCode bool, period uint32) {
for {
util.Println(agentConfig.Debug, "正在更新本地缓存IP信息")
var primaryIP, secondaryIP geoIP
if useIPv6CountryCode {
primaryIP = fetchGeoIP(geoIPApiList, true)
secondaryIP = fetchGeoIP(geoIPApiList, false)
} else {
primaryIP = fetchGeoIP(geoIPApiList, false)
secondaryIP = fetchGeoIP(geoIPApiList, true)
}
ipv4 := fetchGeoIP(geoIPApiList, false)
ipv6 := fetchGeoIP(geoIPApiList, true)
if primaryIP.IP == "" && secondaryIP.IP == "" {
if ipv4.IP == "" && ipv6.IP == "" {
if period > 60 {
time.Sleep(time.Minute)
} else {
@ -75,21 +68,17 @@ func UpdateIP(useIPv6CountryCode bool, period uint32) {
}
continue
}
if primaryIP.IP == "" || secondaryIP.IP == "" {
CachedIP = fmt.Sprintf("%s%s", primaryIP.IP, secondaryIP.IP)
if ipv4.IP == "" || ipv6.IP == "" {
CachedIP = fmt.Sprintf("%s%s", ipv4.IP, ipv6.IP)
} else {
CachedIP = fmt.Sprintf("%s/%s", primaryIP.IP, secondaryIP.IP)
CachedIP = fmt.Sprintf("%s/%s", ipv4.IP, ipv6.IP)
}
if primaryIP.CountryCode != "" {
CachedCountry = primaryIP.CountryCode
} else if secondaryIP.CountryCode != "" {
CachedCountry = secondaryIP.CountryCode
if ipv4.CountryCode != "" {
CachedCountry = ipv4.CountryCode
}
select {
case Sync <- true:
default:
if ipv6.CountryCode != "" && (useIPv6CountryCode || CachedCountry == "") {
CachedCountry = ipv6.CountryCode
}
time.Sleep(time.Second * time.Duration(period))

8
pkg/pty/pty_windows.go
View File

@ -58,9 +58,8 @@ func VersionCheck() bool {
}
return version >= 17763
} else {
return false
}
return false
}
func DownloadDependency() {
@ -131,10 +130,9 @@ func Start() (Pty, error) {
if !isWin10 {
tty, err := winpty.OpenDefault(path, shellPath)
return &winPTY{tty: tty}, err
} else {
tty, err := conpty.Start(shellPath, conpty.ConPtyWorkDir(path))
return &conPty{tty: tty}, err
}
tty, err := conpty.Start(shellPath, conpty.ConPtyWorkDir(path))
return &conPty{tty: tty}, err
}
func (w *winPTY) Write(p []byte) (n int, err error) {