管理端全部功能跑通

EdgeHttpDNS/internal/apps/app_cmd.go

@@ -0,0 +1,149 @@
+package apps
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"runtime"
+	"time"
+
+	teaconst "github.com/TeaOSLab/EdgeHttpDNS/internal/const"
+	"github.com/iwind/TeaGo/maps"
+	"github.com/iwind/gosock/pkg/gosock"
+)
+
+type AppCmd struct {
+	product    string
+	version    string
+	usage      string
+	directives map[string]func()
+	sock       *gosock.Sock
+}
+
+func NewAppCmd() *AppCmd {
+	return &AppCmd{
+		directives: map[string]func(){},
+		sock:       gosock.NewTmpSock(teaconst.ProcessName),
+	}
+}
+
+func (a *AppCmd) Product(product string) *AppCmd {
+	a.product = product
+	return a
+}
+
+func (a *AppCmd) Version(version string) *AppCmd {
+	a.version = version
+	return a
+}
+
+func (a *AppCmd) Usage(usage string) *AppCmd {
+	a.usage = usage
+	return a
+}
+
+func (a *AppCmd) On(arg string, callback func()) {
+	a.directives[arg] = callback
+}
+
+func (a *AppCmd) Run(main func()) {
+	args := os.Args[1:]
+	if len(args) == 0 {
+		main()
+		return
+	}
+
+	switch args[0] {
+	case "-v", "version", "-version", "--version":
+		fmt.Println(a.product+" v"+a.version, "(build:", runtimeString()+")")
+		return
+	case "help", "-h", "--help":
+		fmt.Println(a.product + " v" + a.version)
+		fmt.Println("Usage:")
+		fmt.Println("  " + a.usage)
+		return
+	case "start":
+		a.runDirective("start:before")
+		a.runStart()
+		return
+	case "stop":
+		a.runStop()
+		return
+	case "restart":
+		a.runStop()
+		time.Sleep(1 * time.Second)
+		a.runDirective("start:before")
+		a.runStart()
+		return
+	case "status":
+		a.runStatus()
+		return
+	default:
+		if callback, ok := a.directives[args[0]]; ok {
+			callback()
+			return
+		}
+		fmt.Println("unknown command '" + args[0] + "'")
+	}
+}
+
+func (a *AppCmd) runStart() {
+	pid := a.getPID()
+	if pid > 0 {
+		fmt.Println(a.product+" already started, pid:", pid)
+		return
+	}
+
+	cmd := exec.Command(os.Args[0])
+	cmd.Env = append(os.Environ(), "EdgeBackground=on")
+	err := cmd.Start()
+	if err != nil {
+		fmt.Println(a.product+" start failed:", err.Error())
+		return
+	}
+
+	fmt.Println(a.product+" started, pid:", cmd.Process.Pid)
+}
+
+func (a *AppCmd) runStop() {
+	pid := a.getPID()
+	if pid == 0 {
+		fmt.Println(a.product + " not started")
+		return
+	}
+
+	_, _ = a.sock.Send(&gosock.Command{Code: "stop"})
+	fmt.Println(a.product+" stopped, pid:", pid)
+}
+
+func (a *AppCmd) runStatus() {
+	pid := a.getPID()
+	if pid == 0 {
+		fmt.Println(a.product + " not started")
+		return
+	}
+	fmt.Println(a.product+" is running, pid:", pid)
+}
+
+func (a *AppCmd) runDirective(name string) {
+	if callback, ok := a.directives[name]; ok && callback != nil {
+		callback()
+	}
+}
+
+func (a *AppCmd) getPID() int {
+	if !a.sock.IsListening() {
+		return 0
+	}
+
+	reply, err := a.sock.Send(&gosock.Command{Code: "pid"})
+	if err != nil {
+		return 0
+	}
+
+	return maps.NewMap(reply.Params).GetInt("pid")
+}
+
+func runtimeString() string {
+	return runtime.GOOS + "/" + runtime.GOARCH
+}