✨ v0.9.22 WebSSH

cmd/agent/main.go

@@ -14,20 +14,18 @@ import (
 	"net/url"
 	"os"
 	"os/exec"
-	"runtime"
-	"syscall"
 	"time"
-	"unsafe"
 
 	"github.com/blang/semver"
 	"github.com/genkiroid/cert"
 	"github.com/go-ping/ping"
 	"github.com/gorilla/websocket"
-	"github.com/kr/pty"
 	"github.com/p14yground/go-github-selfupdate/selfupdate"
 	"google.golang.org/grpc"
 
 	"github.com/naiba/nezha/cmd/agent/monitor"
+	"github.com/naiba/nezha/cmd/agent/processgroup"
+	"github.com/naiba/nezha/cmd/agent/pty"
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/pkg/utils"
 	pb "github.com/naiba/nezha/proto"
@@ -97,6 +95,8 @@ func run() {
 		ClientSecret: clientSecret,
 	}
 
+	go pty.DownloadDependency()
+
 	// 上报服务器信息
 	go reportState()
 	// 更新IP信息
@@ -311,7 +311,7 @@ func handleCommandTask(task *pb.Task, result *pb.TaskResult) {
 	startedAt := time.Now()
 	var cmd *exec.Cmd
 	var endCh = make(chan struct{})
-	pg, err := utils.NewProcessExitGroup()
+	pg, err := processgroup.NewProcessExitGroup()
 	if err != nil {
 		// 进程组创建失败，直接退出
 		result.Data = err.Error()
@@ -345,6 +345,11 @@ func handleCommandTask(task *pb.Task, result *pb.TaskResult) {
 	result.Delay = float32(time.Since(startedAt).Seconds())
 }
 
+type WindowSize struct {
+	Cols uint32
+	Rows uint32
+}
+
 func handleTerminalTask(task *pb.Task) {
 	var terminal model.TerminalTask
 	err := json.Unmarshal([]byte(task.GetData()), &terminal)
@@ -365,36 +370,18 @@ func handleTerminalTask(task *pb.Task) {
 	}
 	defer conn.Close()
 
-	var cmd *exec.Cmd
-	var shellPath string
-	if runtime.GOOS == "windows" {
-		shellPath, err = exec.LookPath("powershell.exe")
-		if err != nil || shellPath == "" {
-			shellPath = "cmd.exe"
-		}
-	} else {
-		shellPath = os.Getenv("SHELL")
-		if shellPath == "" {
-			shellPath = "sh"
-		}
-	}
-	cmd = exec.Command(shellPath)
-	cmd.Env = append(os.Environ(), "TERM=xterm")
-
-	tty, err := pty.Start(cmd)
+	tty, err := pty.Start()
 	if err != nil {
 		println("Terminal pty.Start失败：", err)
 		return
 	}
 
 	defer func() {
-		cmd.Process.Kill()
-		cmd.Process.Wait()
 		tty.Close()
 		conn.Close()
 		println("terminal exit", terminal.Session)
 	}()
-	println("terminal init", terminal.Session, shellPath)
+	println("terminal init", terminal.Session)
 
 	go func() {
 		for {
@@ -434,17 +421,12 @@ func handleTerminalTask(task *pb.Task) {
 			io.Copy(tty, reader)
 		case 1:
 			decoder := json.NewDecoder(reader)
-			resizeMessage := windowSize{}
+			var resizeMessage WindowSize
 			err := decoder.Decode(&resizeMessage)
 			if err != nil {
 				continue
 			}
-			syscall.Syscall(
-				syscall.SYS_IOCTL,
-				tty.Fd(),
-				syscall.TIOCSWINSZ,
-				uintptr(unsafe.Pointer(&resizeMessage)),
-			)
+			tty.Setsize(resizeMessage.Cols, resizeMessage.Rows)
 		}
 	}
 }

cmd/agent/processgroup/process_group.go

@@ -0,0 +1,31 @@
+// +build !windows
+
+package processgroup
+
+import (
+	"os/exec"
+	"syscall"
+)
+
+type ProcessExitGroup struct {
+	cmds []*exec.Cmd
+}
+
+func NewProcessExitGroup() (ProcessExitGroup, error) {
+	return ProcessExitGroup{}, nil
+}
+
+func (g *ProcessExitGroup) Dispose() error {
+	for _, c := range g.cmds {
+		if err := syscall.Kill(-c.Process.Pid, syscall.SIGKILL); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (g *ProcessExitGroup) AddProcess(cmd *exec.Cmd) error {
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+	g.cmds = append(g.cmds, cmd)
+	return nil
+}

cmd/agent/processgroup/process_group_windows.go

@@ -0,0 +1,30 @@
+// +build windows
+
+package processgroup
+
+import (
+	"fmt"
+	"os/exec"
+)
+
+type ProcessExitGroup struct {
+	cmds []*exec.Cmd
+}
+
+func NewProcessExitGroup() (ProcessExitGroup, error) {
+	return ProcessExitGroup{}, nil
+}
+
+func (g *ProcessExitGroup) Dispose() error {
+	for _, c := range g.cmds {
+		if err := exec.Command("taskkill", "/F", "/T", "/PID", fmt.Sprint(c.Process.Pid)).Run(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (g *ProcessExitGroup) AddProcess(cmd *exec.Cmd) error {
+	g.cmds = append(g.cmds, cmd)
+	return nil
+}

cmd/agent/pty/pty.go

@@ -0,0 +1,52 @@
+//go:build !windows
+//+build !windows
+
+package pty
+
+import (
+	"os"
+	"os/exec"
+
+	opty "github.com/creack/pty"
+)
+
+type Pty struct {
+	tty *os.File
+	cmd *exec.Cmd
+}
+
+func DownloadDependency() {
+}
+
+func Start() (*Pty, error) {
+	shellPath := os.Getenv("SHELL")
+	if shellPath == "" {
+		shellPath = "sh"
+	}
+	cmd := exec.Command(shellPath)
+	cmd.Env = append(os.Environ(), "TERM=xterm")
+	tty, err := opty.Start(cmd)
+	return &Pty{tty: tty, cmd: cmd}, 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 opty.Setsize(pty.tty, &opty.Winsize{
+		Cols: uint16(cols),
+		Rows: uint16(rows),
+	})
+}
+
+func (pty *Pty) Close() error {
+	if err := pty.tty.Close(); err != nil {
+		return err
+	}
+	return pty.cmd.Process.Kill()
+}

cmd/agent/pty/pty_windows.go

@@ -0,0 +1,96 @@
+// go:build windows
+// +build windows
+
+package pty
+
+import (
+	"fmt"
+	"io/ioutil"
+	"log"
+	"net/http"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+
+	"github.com/artdarek/go-unzip"
+	"github.com/iamacarpet/go-winpty"
+)
+
+type Pty struct {
+	tty *winpty.WinPTY
+}
+
+func DownloadDependency() {
+	resp, err := http.Get("https://dn-dao-github-mirror.daocloud.io/rprichard/winpty/releases/download/0.4.3/winpty-0.4.3-msvc2015.zip")
+	if err != nil {
+		log.Println("wintty 下载失败", err)
+		return
+	}
+	defer resp.Body.Close()
+	content, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		log.Println("wintty 下载失败", err)
+		return
+	}
+	if err := ioutil.WriteFile("./wintty.zip", content, os.FileMode(0777)); err != nil {
+		log.Println("wintty 写入失败", err)
+		return
+	}
+	if err := unzip.New("./wintty.zip", "./wintty").Extract(); err != nil {
+		fmt.Println("wintty 解压失败", err)
+		return
+	}
+	arch := "x64"
+	if runtime.GOARCH != "amd64" {
+		arch = "ia32"
+	}
+	executablePath, err := getExecutableFilePath()
+	if err != nil {
+		fmt.Println("wintty 获取文件路径失败", err)
+		return
+	}
+	os.Rename("./wintty/"+arch+"/bin/winpty-agent.exe", filepath.Join(executablePath, "winpty-agent.exe"))
+	os.Rename("./wintty/"+arch+"/bin/winpty.dll", filepath.Join(executablePath, "winpty.dll"))
+	os.RemoveAll("./wintty")
+	os.RemoveAll("./wintty.zip")
+}
+
+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 := winpty.Open(path, shellPath)
+	return &Pty{tty: tty}, err
+}
+
+func (pty *Pty) Write(p []byte) (n int, err error) {
+	return pty.tty.StdIn.Read(p)
+}
+
+func (pty *Pty) Read(p []byte) (n int, err error) {
+	return pty.tty.StdOut.Read(p)
+}
+
+func (pty *Pty) Setsize(cols, rows uint32) error {
+	pty.tty.SetSize(cols, rows)
+	return nil
+}
+
+func (pty *Pty) Close() error {
+	pty.tty.Close()
+	return nil
+}

cmd/playground/main.go

@@ -1,191 +1,7 @@
 package main
 
-import (
-	"context"
-	"crypto/tls"
-	"fmt"
-	"io"
-	"log"
-	"net"
-	"net/http"
-	"os"
-	"os/exec"
-	"strings"
-	"time"
-
-	"github.com/genkiroid/cert"
-	"github.com/go-ping/ping"
-	"github.com/shirou/gopsutil/v3/cpu"
-	"github.com/shirou/gopsutil/v3/disk"
-	"github.com/shirou/gopsutil/v3/host"
-
-	"github.com/naiba/nezha/pkg/utils"
-)
+import "github.com/naiba/nezha/cmd/agent/pty"
 
 func main() {
-	// icmp()
-	// tcpping()
-	// httpWithSSLInfo()
-	// sysinfo()
-	// cmdExec()
-	// resolveIP("ipapi.co", true)
-	// resolveIP("ipapi.co", false)
-	log.Println(exec.LookPath("powershell.exe"))
-	defaultShell := os.Getenv("SHELL")
-	if defaultShell == "" {
-		defaultShell = "sh"
-	}
-	cmd := exec.Command(defaultShell)
-	cmd.Stdin = os.Stdin
-	stdoutReader, err := cmd.StdoutPipe()
-	if err != nil {
-		println("Terminal StdoutPipe:", err)
-		return
-	}
-	stderrReader, err := cmd.StderrPipe()
-	if err != nil {
-		println("Terminal StderrPipe: ", err)
-		return
-	}
-
-	readers := []io.Reader{stdoutReader, stderrReader}
-	for i := 0; i < len(readers); i++ {
-		go func(j int) {
-			data := make([]byte, 2048)
-			for {
-				count, err := readers[j].Read(data)
-				if err != nil {
-					panic(err)
-				}
-				os.Stdout.Write(data[:count])
-			}
-		}(i)
-	}
-
-	cmd.Run()
-}
-
-func resolveIP(addr string, ipv6 bool) {
-	url := strings.Split(addr, ":")
-
-	dnsServers := []string{"2606:4700:4700::1001", "2001:4860:4860::8844", "2400:3200::1", "2400:3200:baba::1"}
-	if !ipv6 {
-		dnsServers = []string{"1.0.0.1", "8.8.4.4", "223.5.5.5", "223.6.6.6"}
-	}
-
-	log.Println(net.LookupIP(url[0]))
-	for i := 0; i < len(dnsServers); i++ {
-		dnsServer := dnsServers[i]
-		if ipv6 {
-			dnsServer = "[" + dnsServer + "]"
-		}
-		r := &net.Resolver{
-			PreferGo: true,
-			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
-				d := net.Dialer{
-					Timeout: time.Second * 10,
-				}
-				return d.DialContext(ctx, "udp", dnsServer+":53")
-			},
-		}
-		log.Println(r.LookupIP(context.Background(), "ip", url[0]))
-	}
-}
-
-func tcpping() {
-	start := time.Now()
-	conn, err := net.DialTimeout("tcp", "example.com:80", time.Second*10)
-	if err != nil {
-		panic(err)
-	}
-	conn.Write([]byte("ping\n"))
-	conn.Close()
-	fmt.Println(time.Since(start).Microseconds(), float32(time.Since(start).Microseconds())/1000.0)
-}
-
-func sysinfo() {
-	hi, _ := host.Info()
-	var cpuType string
-	if hi.VirtualizationSystem != "" {
-		cpuType = "Virtual"
-	} else {
-		cpuType = "Physical"
-	}
-	cpuModelCount := make(map[string]int)
-	ci, _ := cpu.Info()
-	for i := 0; i < len(ci); i++ {
-		cpuModelCount[ci[i].ModelName]++
-	}
-	var cpus []string
-	for model, count := range cpuModelCount {
-		cpus = append(cpus, fmt.Sprintf("%s %d %s Core", model, count, cpuType))
-	}
-	os.Exit(0)
-	// 硬盘信息，不使用的原因是会重复统计 Linux、Mac
-	dparts, _ := disk.Partitions(false)
-	for _, part := range dparts {
-		u, _ := disk.Usage(part.Mountpoint)
-		if u != nil {
-			log.Printf("%s %d %d", part.Device, u.Total, u.Used)
-		}
-	}
-}
-
-func httpWithSSLInfo() {
-	// 跳过 SSL 检查
-	transCfg := &http.Transport{
-		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
-	}
-	httpClient := &http.Client{Transport: transCfg, CheckRedirect: func(req *http.Request, via []*http.Request) error {
-		return http.ErrUseLastResponse
-	}}
-	url := "https://ops.naibahq.com"
-	resp, err := httpClient.Get(url)
-	fmt.Println(err, resp)
-	// SSL 证书信息获取
-	c := cert.NewCert(url[8:])
-	fmt.Println(c.Error)
-}
-
-func icmp() {
-	pinger, err := ping.NewPinger("10.10.10.2")
-	if err != nil {
-		panic(err) // Blocks until finished.
-	}
-	pinger.Count = 3000
-	pinger.Timeout = 10 * time.Second
-	if err = pinger.Run(); err != nil {
-		panic(err)
-	}
-	fmt.Println(pinger.PacketsRecv, float32(pinger.Statistics().AvgRtt.Microseconds())/1000.0)
-}
-
-func cmdExec() {
-	execFrom, err := os.Getwd()
-	if err != nil {
-		panic(err)
-	}
-	var cmd *exec.Cmd
-	pg, err := utils.NewProcessExitGroup()
-	if err != nil {
-		panic(err)
-	}
-	if utils.IsWindows() {
-		cmd = exec.Command("cmd", "/c", os.Args[1])
-		// cmd = exec.Command("cmd", "/c", execFrom+"/cmd/playground/example.sh hello asd")
-	} else {
-		cmd = exec.Command("sh", "-c", execFrom+`/cmd/playground/example.sh hello && \
-			echo world!`)
-	}
-	pg.AddProcess(cmd)
-	go func() {
-		time.Sleep(time.Second * 10)
-		if err = pg.Dispose(); err != nil {
-			panic(err)
-		}
-		fmt.Println("killed")
-	}()
-	output, err := cmd.Output()
-	log.Println("output:", string(output))
-	log.Println("err:", err)
+	pty.DownloadDependency()
 }