EdgeNode/internal/apps/app_cmd.go

package apps

import (
	"fmt"
	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
	"github.com/iwind/TeaGo/logs"
	"github.com/iwind/TeaGo/maps"
	"github.com/iwind/TeaGo/types"
	"github.com/iwind/gosock/pkg/gosock"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// AppCmd App命令帮助
type AppCmd struct {
	product       string
	version       string
	usages        []string
	options       []*CommandHelpOption
	appendStrings []string

	directives []*Directive

	sock *gosock.Sock
}

func NewAppCmd() *AppCmd {
	return &AppCmd{
		sock: gosock.NewTmpSock(teaconst.ProcessName),
	}
}

type CommandHelpOption struct {
	Code        string
	Description string
}

// Product 产品
func (this *AppCmd) Product(product string) *AppCmd {
	this.product = product
	return this
}

// Version 版本
func (this *AppCmd) Version(version string) *AppCmd {
	this.version = version
	return this
}

// Usage 使用方法
func (this *AppCmd) Usage(usage string) *AppCmd {
	this.usages = append(this.usages, usage)
	return this
}

// Option 选项
func (this *AppCmd) Option(code string, description string) *AppCmd {
	this.options = append(this.options, &CommandHelpOption{
		Code:        code,
		Description: description,
	})
	return this
}

// Append 附加内容
func (this *AppCmd) Append(appendString string) *AppCmd {
	this.appendStrings = append(this.appendStrings, appendString)
	return this
}

// Print 打印
func (this *AppCmd) Print() {
	fmt.Println(this.product + " v" + this.version)

	fmt.Println("Usage:")
	for _, usage := range this.usages {
		fmt.Println("   " + usage)
	}

	if len(this.options) > 0 {
		fmt.Println("")
		fmt.Println("Options:")

		spaces := 20
		max := 40
		for _, option := range this.options {
			l := len(option.Code)
			if l < max && l > spaces {
				spaces = l + 4
			}
		}

		for _, option := range this.options {
			if len(option.Code) > max {
				fmt.Println("")
				fmt.Println("  " + option.Code)
				option.Code = ""
			}

			fmt.Printf("  %-"+strconv.Itoa(spaces)+"s%s\n", option.Code, ": "+option.Description)
		}
	}

	if len(this.appendStrings) > 0 {
		fmt.Println("")
		for _, s := range this.appendStrings {
			fmt.Println(s)
		}
	}
}

// On 添加指令
func (this *AppCmd) On(arg string, callback func()) {
	this.directives = append(this.directives, &Directive{
		Arg:      arg,
		Callback: callback,
	})
}

// Run 运行
func (this *AppCmd) Run(main func()) {
	// 获取参数
	args := os.Args[1:]
	if len(args) > 0 {
		switch args[0] {
		case "-v", "version", "-version", "--version":
			this.runVersion()
			return
		case "?", "help", "-help", "h", "-h":
			this.runHelp()
			return
		case "start":
			this.runStart()
			return
		case "stop":
			this.runStop()
			return
		case "restart":
			this.runRestart()
			return
		case "status":
			this.runStatus()
			return
		}

		// 查找指令
		for _, directive := range this.directives {
			if directive.Arg == args[0] {
				directive.Callback()
				return
			}
		}

		fmt.Println("unknown command '" + args[0] + "'")

		return
	}

	// 日志
	writer := new(LogWriter)
	writer.Init()
	logs.SetWriter(writer)

	// 运行主函数
	main()
}

// 版本号
func (this *AppCmd) runVersion() {
	fmt.Println(this.product+" v"+this.version, "(build: "+runtime.Version(), runtime.GOOS, runtime.GOARCH, teaconst.Tag+")")
}

// 帮助
func (this *AppCmd) runHelp() {
	this.Print()
}

// 启动
func (this *AppCmd) runStart() {
	var pid = this.getPID()
	if pid > 0 {
		fmt.Println(this.product+" already started, pid:", pid)
		return
	}

	_ = os.Setenv("EdgeBackground", "on")

	cmd := exec.Command(os.Args[0])
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Foreground: false,
		Setsid:     true,
	}

	err := cmd.Start()
	if err != nil {
		fmt.Println(this.product+"  start failed:", err.Error())
		return
	}

	fmt.Println(this.product+" started ok, pid:", cmd.Process.Pid)
}

// 停止
func (this *AppCmd) runStop() {
	var pid = this.getPID()
	if pid == 0 {
		fmt.Println(this.product + " not started yet")
		return
	}

	// 从systemd中停止
	if runtime.GOOS == "linux" {
		systemctl, _ := exec.LookPath("systemctl")
		if len(systemctl) > 0 {
			go func() {
				// 有可能会长时间执行，这里不阻塞进程
				_ = exec.Command(systemctl, "stop", teaconst.SystemdServiceName).Run()
			}()
		}
	}

	// 如果仍在运行，则发送停止指令
	_, _ = this.sock.SendTimeout(&gosock.Command{Code: "stop"}, 1*time.Second)

	fmt.Println(this.product+" stopped ok, pid:", types.String(pid))
}

// 重启
func (this *AppCmd) runRestart() {
	this.runStop()
	time.Sleep(1 * time.Second)
	this.runStart()
}

// 状态
func (this *AppCmd) runStatus() {
	var pid = this.getPID()
	if pid == 0 {
		fmt.Println(this.product + " not started yet")
		return
	}

	fmt.Println(this.product + " is running, pid: " + types.String(pid))
}

// 获取当前的PID
func (this *AppCmd) getPID() int {
	if !this.sock.IsListening() {
		return 0
	}

	reply, err := this.sock.Send(&gosock.Command{Code: "pid"})
	if err != nil {
		return 0
	}
	return maps.NewMap(reply.Params).GetInt("pid")
}

// ParseOptions 分析参数中的选项
func (this *AppCmd) ParseOptions(args []string) map[string][]string {
	var result = map[string][]string{}
	for _, arg := range args {
		var pieces = strings.SplitN(arg, "=", 2)
		var key = strings.TrimLeft(pieces[0], "- ")
		key = strings.TrimSpace(key)
		var value = ""
		if len(pieces) == 2 {
			value = strings.TrimSpace(pieces[1])
		}
		result[key] = append(result[key], value)
	}
	return result
}