rtu_linux_modules/sshd/client/client.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/denisbrodbeck/machineid"
	"github.com/google/uuid"
	"github.com/peterh/liner"
)

const MODULE_NAME = "YFKJ_SSH_CLIENT"

var (
	coupler               *MQTTCoupler
	Version               = "1.0.0.1"
	BuildTime             = "unknown"
	ErrBrokerAddressEmpty = errors.New("mqtt server address is empty")
	ErrIMEINotAvailable   = errors.New("device imei is not available")
)

func main() {
	if IsArgsParam("-h") {
		help()
		return
	}

	if IsArgsParam("-v") {
		fmt.Println("程序版本:", Version, "\n构建时间:", BuildTime)
		return
	}

	devIMEI := GetArgsParamStr("-c", "")
	if devIMEI == "" {
		help()
		return
	}

	if err := loadAppConfig(); err != nil {
		fmt.Printf("[%s] 错误: %v!!\n", MODULE_NAME, err)
		return
	}
	if CfgServers.MQTTSrv.Address == "" {
		fmt.Printf("[%s] 错误: %v!!\n", MODULE_NAME, ErrBrokerAddressEmpty)
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	coupler = &MQTTCoupler{
		ctx:         ctx,
		cancel:      cancel,
		broker:      CfgServers.MQTTSrv.Address,
		username:    CfgServers.MQTTSrv.Username,
		password:    CfgServers.MQTTSrv.Password,
		clientID:    uuid.New().String(),
		imei:        devIMEI,
		cwd:         "/",
		interrupted: make(chan struct{}, 1),
	}

	id, err := machineid.ID()
	if err == nil {
		coupler.machineID = id
	}

	if err := coupler.init2(); err != nil {
		fmt.Printf("[%s] 错误: %v\n!!", MODULE_NAME, err)
		return
	}

	var pingState, pingErrorPrint atomic.Bool
	pingErrorPrint.Store(true)
	heartbeatLoop(&pingState, &pingErrorPrint) /// ping
	for {
		if pingState.Load() { //// 等待成功连接上目标设备卍
			break
		}
		fmt.Printf("[%s] 正在尝试连接设备...\n", MODULE_NAME)
		time.Sleep(1 * time.Second)
	}
	pingErrorPrint.Store(false)

	term(&pingState) /////////////////// 启动终端模拟器卍
}

func term(pingState *atomic.Bool) {
	var executing atomic.Bool   // 是否有正在执行中的命令
	var interrupted atomic.Bool // 用户是否按键取消了命令

	interruptLoop(pingState, &executing, &interrupted) // Ctrl+C卍

	printWelcome(pingState)

	line := liner.NewLiner()
	defer line.Close()
	line.SetCtrlCAborts(false)
	line.SetTabCompletionStyle(liner.TabCircular)

	historyFile := "history.txt"
	var history []string

	if f, err := os.Open(historyFile); err == nil {
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			cmd := strings.TrimSpace(scanner.Text())
			if cmd != "" {
				history = append(history, cmd)
				line.AppendHistory(cmd)
			}
		}
		f.Close()
	}

	defer func() {
		if f, err := os.Create(historyFile); err == nil {
			line.WriteHistory(f)
			f.Close()
		}
	}()

	line.SetCompleter(func(input string) []string { // 补全
		var matches []string
		for _, cmd := range history {
			if strings.HasPrefix(cmd, input) {
				matches = append(matches, cmd)
			}
		}
		return matches
	})

	executing.Store(true)
	result, err := coupler.exec("pwd")
	executing.Store(false)
	if err == nil && result.Cwd != "" { /// 在进入前对齐路径
		coupler.cwd = result.Cwd
	}

	for {
		if !pingState.Load() {
			fmt.Printf("[%s] 目标设备连接丢失!!\n", MODULE_NAME)
			break
		}

		if interrupted.Swap(false) {
			fmt.Println("^C")
		}

		prompt := fmt.Sprintf("root@%s:%s# ", coupler.imei, coupler.cwd)
		input, err := line.Prompt(prompt) /// 等待用户输入指令
		if err == nil {
			goto inputOK
		}

		if err == io.EOF { ////////////////// Ctrl+D 按键处理
			_, _ = coupler.quit()
			fmt.Println("bye")
			break
		}

		fmt.Println(err)
		continue

	inputOK:
		input = strings.TrimSpace(input)
		if input == "" {
			continue
		}

		line.AppendHistory(input) ///// 保存用户输入的历史命令

		history = append(history, input) ///// 本次也立刻生效

		if input == "exit" || input == "quit" {
			_, _ = coupler.quit()
			fmt.Println("bye")
			break
		}

		if input == "clear" {
			clearScreen()
			continue
		}

		executing.Store(true)
		result, err := coupler.exec(input)
		executing.Store(false)

		if err != nil {
			fmt.Printf("%v\n", err)
			continue
		}

		if result.Stdout != "" {
			fmt.Print(result.Stdout)
		}
		if result.Stderr != "" {
			fmt.Fprintln(os.Stderr, result.Stderr)
		}
		if result.ExitCode == 124 {
			fmt.Println("command timeout")
		}

		if result.Cwd != "" {
			coupler.cwd = result.Cwd
		}
	}
}

func help() {
	h := `
-h 显示帮助提示
-v 当前程序版本
-c 连接目标设备(IMEI), 例如: -c 869523059113051
`

	fmt.Println(h)
}

func interruptLoop(pingState, executing, interrupted *atomic.Bool) {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT)
	go func() {
		for range sigCh {
			interrupted.Store(true)
			if executing.Load() {
				select {
				case coupler.interrupted <- struct{}{}:
				default:
				}
				if pingState.Load() {
					go func() {
						_, _ = coupler.stop()
					}()
				} // end if2
			} //// end if1
		} ////// end for
	}()
}

func heartbeatLoop(pingState, pingErrorPrint *atomic.Bool) {
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()
		pingFailCount := 0
		pong := ""
		for range ticker.C {
			resp, err := coupler.ping()
			if err == nil && resp.Error == nil && resp.Result != nil &&
				json.Unmarshal(resp.Result, &pong) == nil && pong == "pong" {
				pingState.Store(true)
				pingFailCount = 0
			} else {
				if pingErrorPrint.Load() { /// 打印的错误信息, 提升交互使用体验
					errMsg := ""
					if err != nil {
						errMsg = err.Error()
					} else if resp.Error != nil {
						errMsg = resp.Error.Message
					} else {
						errMsg = string(resp.Result)
					}
					fmt.Printf("[%s] 错误: 设备未响应(%s)!!\n", MODULE_NAME, errMsg)
				}
				pingFailCount++
				if pingFailCount >= 3 { ///// 连续3次ping失败, 可以认为设备离线
					pingState.Store(false)
				}
			} // end if
		} // end for
	}()
}

func printWelcome(pingState *atomic.Bool) {
	welcome := `
 ____   _   _        _____ _      _   _  _   _ 
|  _ \ | | | |      / ____| |    | | | || \ | |
| |_) || |_| | ___ | |    | | ___| | | ||  \| |
|  _ < |  _  |/ _ \| |    | |/ _ \ | | || . _ |
| |_) || | | |  __/| |____| |  __/ |_| || |\  |
|____/ |_| |_|\___| \_____|_|\___|\___/ |_| \_|  
================================================
      欢迎使用云飞科技RTU设备远程运维终端
提示: 输入"quit"命令或按下"Ctrl+D"键, 可退出终端
================================================
`
	fmt.Println(welcome)
	fmt.Printf("客户端ID: %s\n", coupler.machineID)
	state := ""
	if pingState.Load() {
		state = "√ 已连接"
	} else {
		state = "x 未连接"
	}
	fmt.Printf("设备状态: %s\n\n", state)
}

func clearScreen() {
	if runtime.GOOS == "windows" {
		cmd := exec.Command("cmd", "/c", "cls")
		cmd.Stdout = os.Stdout
		_ = cmd.Run()
	} else {
		cmd := exec.Command("clear")
		cmd.Stdout = os.Stdout
		_ = cmd.Run()
	}
}