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

Switch to ConPty on Windows 10 & support windows arm64 platform (#20)

Browse Source
This commit is contained in:
UUBulb 2024-05-21 18:40:08 +08:00 committed by GitHub
parent 060d13a759
commit f882ca3498
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 296 additions and 146 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.goreleaser.yml
View File

@ -25,8 +25,6 @@ builds:
ignore:
- goos: windows
goarch: arm
- goos: windows
goarch: arm64
main: ./cmd/agent
binary: nezha-agent
- id: darwin-amd64

9
cmd/agent/edit.go
View File

@ -1,16 +1,15 @@
package main
import (
"fmt"
"strings"
"net"
"errors"
"fmt"
"net"
"strings"
"github.com/spf13/cobra"
"github.com/AlecAivazis/survey/v2"
"github.com/shirou/gopsutil/v3/disk"
psnet "github.com/shirou/gopsutil/v3/net"
"github.com/spf13/cobra"
)
var editCmd = &cobra.Command{

109
cmd/agent/main.go
View File

@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"io"
"log"
"net"
"net/http"
"net/url"
@ -22,6 +23,7 @@ import (
"github.com/go-ping/ping"
"github.com/gorilla/websocket"
"github.com/nezhahq/go-github-selfupdate/selfupdate"
"github.com/nezhahq/service"
"github.com/quic-go/quic-go/http3"
"github.com/shirou/gopsutil/v3/host"
"github.com/spf13/cobra"
@ -52,17 +54,27 @@ type AgentCliParam struct {
IPReportPeriod uint32 // 上报IP间隔
}
type program struct {
exit chan struct{}
service service.Service
}
var (
version string
arch string
client pb.NezhaServiceClient
inited bool
logger service.Logger
)
var agentCmd = &cobra.Command{
Use: "agent",
Run: func(cmd *cobra.Command, args []string) {
run()
if runtime.GOOS == "darwin" {
run() // macOS launchctl 如使用 runService 则无法启动，原因未知
} else {
runService("")
}
},
PreRun: preRun,
PersistentPreRun: persistPreRun,
@ -196,6 +208,33 @@ func preRun(cmd *cobra.Command, args []string) {
}
}
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()
return
}
func run() {
auth := model.AuthHandler{
ClientSecret: agentCliParam.ClientSecret,
@ -274,6 +313,74 @@ func run() {
}
}
func runService(action string) {
var tlsoption string
dir, err := os.Getwd()
if err != nil {
println("获取当前工作目录时出错: ", err)
return
}
if agentCliParam.TLS {
tlsoption = "--tls"
}
svcConfig := &service.Config{
Name: "nezha-agent",
DisplayName: "Nezha Agent",
Description: "哪吒探针监控端",
Arguments: []string{
"-s", agentCliParam.Server,
"-p", agentCliParam.ClientSecret,
tlsoption,
},
WorkingDirectory: dir,
}
prg := &program{
exit: make(chan struct{}),
}
s, err := service.New(prg, svcConfig)
if err != nil {
log.Fatal("创建服务时出错: ", err)
}
prg.service = s
errs := make(chan error, 5)
logger, err = s.Logger(errs)
if err != nil {
log.Fatal(err)
}
go func() {
for {
err := <-errs
if err != nil {
log.Print(err)
}
}
}()
if action == "install" {
initName := s.Platform()
log.Println("Init system is:", initName)
}
if len(action) != 0 {
err := service.Control(s, action)
if err != nil {
log.Fatal(err)
}
return
}
err = s.Run()
if err != nil {
logger.Error(err)
}
}
func receiveTasks(tasks pb.NezhaService_RequestTaskClient) error {
var err error
defer println("receiveTasks exit", time.Now(), "=>", err)

82
cmd/agent/service.go
View File

@ -2,22 +2,15 @@ package main
import (
"os"
"log"
"github.com/spf13/cobra"
"github.com/nezhahq/service"
)
type program struct {
exit chan struct{}
service service.Service
}
var serviceCmd = &cobra.Command{
Use: "service <install/uninstall/start/stop/restart>",
Short: "服务与自启动设置",
Args: cobra.ExactArgs(1),
Run: runService,
Run: serviceActions,
PreRun: servicePreRun,
}
@ -39,74 +32,7 @@ func servicePreRun(cmd *cobra.Command, args []string) {
}
}
func (p *program) Start(s service.Service) error {
go p.run()
return nil
}
func (p *program) run() {
defer func() {
if service.Interactive() {
p.Stop(p.service)
} else {
p.service.Stop()
}
}()
run()
return
}
func (p *program) Stop(s service.Service) error {
close(p.exit)
if service.Interactive() {
os.Exit(0)
}
return nil
}
func runService(cmd *cobra.Command, args []string) {
var tlsoption string
mode := args[0]
dir, err := os.Getwd()
if err != nil {
println("获取当前工作目录时出错: ", err)
return
}
if agentCliParam.TLS {
tlsoption = "--tls"
}
svcConfig := &service.Config{
Name: "nezha-agent",
DisplayName: "Nezha Agent",
Description: "哪吒探针监控端",
Arguments: []string{
"-s", agentCliParam.Server,
"-p", agentCliParam.ClientSecret,
tlsoption,
},
WorkingDirectory: dir,
}
prg := &program{
exit: make(chan struct{}),
}
s, err := service.New(prg, svcConfig)
if err != nil {
log.Fatal("创建服务时出错: ", err)
}
if mode == "install" {
initName := s.Platform()
log.Println("Init system is:", initName)
}
err = service.Control(s, mode)
if err != nil {
log.Fatal(err)
}
func serviceActions(cmd *cobra.Command, args []string) {
action := args[0]
runService(action)
}

5
go.mod
View File

@ -5,6 +5,7 @@ go 1.20
require (
github.com/AlecAivazis/survey/v2 v2.3.7
github.com/DaRealFreak/cloudflare-bp-go v1.0.4
github.com/UserExistsError/conpty v0.1.3
github.com/artdarek/go-unzip v1.0.0
github.com/blang/semver v3.5.1+incompatible
github.com/creack/pty v1.1.21
@ -15,7 +16,7 @@ require (
github.com/iamacarpet/go-winpty v1.0.4
github.com/json-iterator/go v1.1.12
github.com/nezhahq/go-github-selfupdate v0.0.0-20240418134522-9d84a13bbf2d
github.com/nezhahq/service v0.0.0-20240518110122-594be3ba1962
github.com/nezhahq/service v0.0.0-20240519060034-90701692f8ef
github.com/quic-go/quic-go v0.40.1
github.com/shirou/gopsutil/v3 v3.24.4
github.com/spf13/cobra v1.8.0
@ -77,7 +78,7 @@ require (
golang.org/x/net v0.21.0 // indirect
golang.org/x/oauth2 v0.17.0 // indirect
golang.org/x/sync v0.6.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/sys v0.20.0 // indirect
golang.org/x/term v0.17.0 // indirect
golang.org/x/text v0.14.0 // indirect
golang.org/x/tools v0.13.0 // indirect

10
go.sum
View File

@ -8,6 +8,8 @@ github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2 h1:+vx7roKuyA63n
github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2/go.mod h1:HBCaDeC1lPdgDeDbhX8XFpy1jqjK0IBG8W5K+xYqA0w=
github.com/PuerkitoBio/goquery v1.8.1 h1:uQxhNlArOIdbrH1tr0UXwdVFgDcZDrZVdcpygAcwmWM=
github.com/PuerkitoBio/goquery v1.8.1/go.mod h1:Q8ICL1kNUJ2sXGoAhPGUdYDJvgQgHzJsnnd3H7Ho5jQ=
github.com/UserExistsError/conpty v0.1.3 h1:YzGQkHAiBBkAihOCO5J2cAnahzb8ePvje2YxG7et1E0=
github.com/UserExistsError/conpty v0.1.3/go.mod h1:PDglKIkX3O/2xVk0MV9a6bCWxRmPVfxqZoTG/5sSd9I=
github.com/andybalholm/cascadia v1.3.1 h1:nhxRkql1kdYCc8Snf7D5/D3spOX+dBgjA6u8x004T2c=
github.com/andybalholm/cascadia v1.3.1/go.mod h1:R4bJ1UQfqADjvDa4P6HZHLh/3OxWWEqc0Sk8XGwHqvA=
github.com/artdarek/go-unzip v1.0.0 h1:Ja9wfhiXyl67z5JT37rWjTSb62KXDP+9jHRkdSREUvg=
@ -101,12 +103,10 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G
github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
github.com/nezhahq/go-github-selfupdate v0.0.0-20240418134522-9d84a13bbf2d h1:gUt6JLTE/HH7qXS5F7SSYrmsj8NXx3i5CaF6mzYw6pA=
github.com/nezhahq/go-github-selfupdate v0.0.0-20240418134522-9d84a13bbf2d/go.mod h1:fOsabb2tjwCe7/kGSsout6oL2cp0sKhOwYQp6fP/Xfg=
github.com/nezhahq/service v0.0.0-20240518043736-9ae0db11e8df h1:V9mIEc9CKm+hFuPmrEjByCScbWZGoWLOUcPjrAis9ew=
github.com/nezhahq/service v0.0.0-20240518043736-9ae0db11e8df/go.mod h1:1CemEvuMOM4B88ckxMZvLT0dehlP5+bqQOYRLMslSdE=
github.com/nezhahq/service v0.0.0-20240518100746-f22fc3c61e5d h1:OmxtpsZfT0wLBgwo0uw4vusU+nEVWljye0vWfRpbFig=
github.com/nezhahq/service v0.0.0-20240518100746-f22fc3c61e5d/go.mod h1:1CemEvuMOM4B88ckxMZvLT0dehlP5+bqQOYRLMslSdE=
github.com/nezhahq/service v0.0.0-20240518110122-594be3ba1962 h1:2zrICSRNedjDMlmDb58B03QHcLvpKIBe931FV7kHvas=
github.com/nezhahq/service v0.0.0-20240518110122-594be3ba1962/go.mod h1:1CemEvuMOM4B88ckxMZvLT0dehlP5+bqQOYRLMslSdE=
github.com/nezhahq/service v0.0.0-20240519060034-90701692f8ef h1:nC+Cxc9MyBFT+snkl5lBUpxGqSGqXv9B2AMHtZOtXzo=
github.com/nezhahq/service v0.0.0-20240519060034-90701692f8ef/go.mod h1:i6zO7Vzuv5+mdaCzHrvAC4U63W59uXmX9n6o7p4PJGk=
github.com/onsi/ginkgo/v2 v2.9.5 h1:+6Hr4uxzP4XIUyAkg61dWBw8lb/gc4/X5luuxN/EC+Q=
github.com/onsi/ginkgo/v2 v2.9.5/go.mod h1:tvAoo1QUJwNEU2ITftXTpR7R1RbCzoZUOs3RonqW57k=
github.com/onsi/gomega v1.27.6 h1:ENqfyGeS5AX/rlXDd/ETokDz93u0YufY1Pgxuy/PvWE=
@ -221,6 +221,8 @@ golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.19.0 h1:q5f1RH2jigJ1MoAWp2KTp3gm5zAGFUTarQZ5U386+4o=
golang.org/x/sys v0.19.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y=
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=

75
pkg/pty/pty_windows.go
View File

@ -1,4 +1,4 @@
//go:build windows
//go:build windows && !arm64
package pty
@ -10,17 +10,50 @@ import (
"os"
"os/exec"
"path/filepath"
"regexp"
"runtime"
"strconv"
"github.com/UserExistsError/conpty"
"github.com/artdarek/go-unzip"
"github.com/iamacarpet/go-winpty"
"github.com/shirou/gopsutil/v3/host"
)
var isWin10 bool
type Pty struct {
tty *winpty.WinPTY
tty interface{}
}
func init() {
isWin10 = VersionCheck()
}
func VersionCheck() bool {
hi, err := host.Info()
if err != nil {
return false
}
re := regexp.MustCompile(`Build (\d+(\.\d+)?)`)
match := re.FindStringSubmatch(hi.PlatformVersion)
if len(match) > 1 {
versionStr := match[1]
version, err := strconv.ParseFloat(versionStr, 64)
if err != nil {
return false
}
return version >= 17763
} else {
return false
}
}
func DownloadDependency() {
if !isWin10 {
executablePath, err := getExecutableFilePath()
if err != nil {
fmt.Println("NEZHA>> wintty 获取文件路径失败", err)
@ -36,7 +69,7 @@ func DownloadDependency() {
return
}
resp, err := http.Get("https://dn-dao-github-mirror.daocloud.io/rprichard/winpty/releases/download/0.4.3/winpty-0.4.3-msvc2015.zip")
resp, err := http.Get("https://github.com/rprichard/winpty/releases/download/0.4.3/winpty-0.4.3-msvc2015.zip")
if err != nil {
log.Println("NEZHA>> wintty 下载失败", err)
return
@ -64,6 +97,7 @@ func DownloadDependency() {
os.Rename("./wintty/"+arch+"/bin/winpty.dll", winptyAgentDll)
os.RemoveAll("./wintty")
os.RemoveAll("./wintty.zip")
}
}
func getExecutableFilePath() (string, error) {
@ -75,6 +109,8 @@ func getExecutableFilePath() (string, error) {
}
func Start() (*Pty, error) {
var tty interface{}
shellPath, err := exec.LookPath("powershell.exe")
if err != nil || shellPath == "" {
shellPath = "cmd.exe"
@ -83,24 +119,47 @@ func Start() (*Pty, error) {
if err != nil {
return nil, err
}
tty, err := winpty.OpenDefault(path, shellPath)
if !isWin10 {
tty, err = winpty.OpenDefault(path, shellPath)
} else {
tty, err = conpty.Start(shellPath, conpty.ConPtyWorkDir(path))
}
return &Pty{tty: tty}, err
}
func (pty *Pty) Write(p []byte) (n int, err error) {
return pty.tty.StdIn.Write(p)
if !isWin10 {
return pty.tty.(*winpty.WinPTY).StdIn.Write(p)
} else {
return pty.tty.(*conpty.ConPty).Write(p)
}
}
func (pty *Pty) Read(p []byte) (n int, err error) {
return pty.tty.StdOut.Read(p)
if !isWin10 {
return pty.tty.(*winpty.WinPTY).StdOut.Read(p)
} else {
return pty.tty.(*conpty.ConPty).Read(p)
}
}
func (pty *Pty) Setsize(cols, rows uint32) error {
pty.tty.SetSize(cols, rows)
if !isWin10 {
pty.tty.(*winpty.WinPTY).SetSize(cols, rows)
return nil
} else {
return pty.tty.(*conpty.ConPty).Resize(int(cols), int(rows))
}
}
func (pty *Pty) Close() error {
pty.tty.Close()
if !isWin10 {
pty.tty.(*winpty.WinPTY).Close()
return nil
} else {
if err := pty.tty.(*conpty.ConPty).Close(); err != nil {
return err
}
return nil
}
}

58
pkg/pty/pty_windowsarm.go Normal file
View File

@ -0,0 +1,58 @@
//go:build windows && arm64
package pty
import (
"os"
"os/exec"
"path/filepath"
"github.com/UserExistsError/conpty"
)
type Pty struct {
tty *conpty.ConPty
}
func DownloadDependency() {
}
func getExecutableFilePath() (string, error) {
ex, err := os.Executable()
if err != nil {
return "", err
}
return filepath.Dir(ex), nil
}
func Start() (*Pty, error) {
shellPath, err := exec.LookPath("powershell.exe")
if err != nil || shellPath == "" {
shellPath = "cmd.exe"
}
path, err := getExecutableFilePath()
if err != nil {
return nil, err
}
tty, err := conpty.Start(shellPath, conpty.ConPtyWorkDir(path))
return &Pty{tty: tty}, err
}
func (pty *Pty) Write(p []byte) (n int, err error) {
return pty.tty.Write(p)
}
func (pty *Pty) Read(p []byte) (n int, err error) {
return pty.tty.Read(p)
}
func (pty *Pty) Setsize(cols, rows uint32) error {
return pty.tty.Resize(int(cols), int(rows))
}
func (pty *Pty) Close() error {
if err := pty.tty.Close(); err != nil {
return err
}
return nil
}