pmt/main.go

package main

import (
	"flag"
	"fmt"
	"github.com/rs/xid"
	"hupeh.vip/pm/cmd"
	"hupeh.vip/pm/db"
	"os"
	"os/exec"
	"strconv"
)

var allowPMTs = map[string]int{
	"start":   1,
	"restart": 1,
	//"list": 1,
	//"shell": 1,
}

func main() {
	var sign string // all
	var pmt string  // all
	var id int      // start,restart
	var daemon bool // start,restart

	flag.StringVar(&sign, "sign", "", "sign token")
	flag.StringVar(&pmt, "pmt", "", "task name")
	flag.IntVar(&id, "id", 0, "app id")
	flag.BoolVar(&daemon, "daemon", false, "daemon mode")
	flag.Parse()

	// 命令行模式
	if v, ok := allowPMTs[pmt]; !ok || v != 1 {
		runCli()
		return
	}

	// 验证签名
	if _, err := xid.FromString(sign); err != nil {
		fmt.Println(cmd.Red("bad sign token"))
		return
	}

	switch pmt {
	// 启动应用
	case "start":
		// 开启守护模式
		if !daemon {
			subProcess(append(os.Args, "-daemon")...)
			fmt.Printf("[*] Daemon running in pid: %d PPID: %d\n", os.Getpid(), os.Getppid())
			os.Exit(0)
		}
		if id <= 0 {
			fmt.Println(cmd.Red("missing app id"))
			return
		}
		// 查询 app
		app, err := db.FindApp(db.WithAppID(uint(id)))
		if err != nil {
			fmt.Println(cmd.Red("failed to find app"))
			fmt.Println("    " + cmd.Red(err.Error()))
			return
		}
		// 创建进程处理器
		proc := db.NewProc(app)
		// 忽略运行中的
		if proc.IsAlive() {
			return
		}
		// 启动失败
		if err = proc.Start(); err != nil {
			fmt.Println(cmd.Yellow(app.Name) + " - " + cmd.Red("start fail"))
			fmt.Println("    " + cmd.Red(err.Error()))
			return
		}
		// 启动成功
		fmt.Println(cmd.Yellow(app.Name) + " - " + cmd.Green("start success"))
	// 重启脚本
	case "restart":
		// 开启守护模式
		if !daemon {
			subProcess(append(os.Args, "-daemon")...)
			fmt.Printf("[*] Daemon running in pid: %d PPID: %d\n", os.Getpid(), os.Getppid())
			os.Exit(0)
		}
		if id <= 0 {
			fmt.Println(cmd.Red("missing app id"))
			return
		}
		// 查询 app
		app, err := db.FindApp(db.WithAppID(uint(id)))
		if err != nil {
			fmt.Println(cmd.Red("failed to find app"))
			fmt.Println("    " + cmd.Red(err.Error()))
			return
		}
		// 创建进程处理器
		proc := db.NewProc(app)
		// 重启失败
		if err = proc.Restart(); err != nil {
			fmt.Println(cmd.Yellow(app.Name) + " - " + cmd.Red("restart fail"))
			fmt.Println("    " + cmd.Red(err.Error()))
			return
		}
		// 重启成功
		fmt.Println(cmd.Yellow(app.Name) + " - " + cmd.Green("restart success"))
	}
}

func runCli() {
	cmd.New("pm").
		Name("pm").
		Description("description").

		// start 命令
		Command("start", "Start program(s) in daemon").
		Argument("[names:...string]", "file/name/id").
		Option("-w, --watch <watch:boolean>", "Watches folder changes").
		Option("-c, --cwd <cwd:string>", "Sets working directory    "+cmd.Yellow("Sets working directory")+" Sets working directory "+cmd.Yellow("Sets working directory")+" Sets working directory "+cmd.Yellow(" Sets working directory")+" Sets working directory "+cmd.Yellow("Sets working directory")+" Sets working directory "+cmd.Yellow("Sets working directory")+" End ...").
		Option("-n, --name", "Sets program name").
		Option("-e, --env", "Sets current environment name").
		Option("-i, --interpreter", "Sets interpreter name").
		Action(func(c *cmd.Command, m map[string]any, a ...any) error {
			var apps []*db.App
			var err error
			if len(a) != 1 {
				apps, err = db.FindApps(db.WithAppStatus(db.StatusStopped))
				if err != nil {
					fmt.Println(cmd.Red("load all apps fail"))
					fmt.Println("    " + cmd.Red(err.Error()))
					return err
				}
			} else {
				scripts := a[0].([]string)
				for _, script := range scripts {
					loads, err := db.ParseInput(script)
					if err != nil {
						fmt.Println(cmd.Red("load all apps fail"))
						fmt.Println("    " + cmd.Red(err.Error()))
						return err
					}
					apps = append(apps, loads...)
				}
			}
			if len(apps) == 0 {
				fmt.Println("No app found")
				return nil
			}
			for _, app := range apps {
				if app.Status == db.StatusRunning {
					fmt.Println(cmd.Yellow(app.Name), cmd.Green("running"))
					continue
				}
				if app.ID == 0 {
					if err := db.SaveApp(app); err != nil {
						fmt.Println(cmd.Yellow(app.Name), cmd.Red("failed to saving"))
						fmt.Println("    " + cmd.Red(err.Error()))
						continue
					}
				}
				// todo options
				proc := subProcess(
					"-id="+strconv.Itoa(int(app.ID)),
					"-pmt=start",
					"-sign="+xid.New().String(),
				)
				if err := proc.Wait(); err != nil {
					fmt.Println(cmd.Yellow(app.Name), cmd.Red("failed"))
					fmt.Println("    ", err.Error())
				} else {
					fmt.Println(cmd.Yellow(app.Name), cmd.Green("success"))
				}
			}
			fmt.Println("over")
			return nil
		}).

		// restart 命令
		Command("restart", "Restarts program(s)").
		Argument("<names:...string>", "file/name/id").
		Action(func(c *cmd.Command, m map[string]any, a ...any) error {
			var apps []*db.App
			if len(a) != 0 {
				c.ShowHelp()
				return nil
			}
			scripts := a[0].([]string)
			for _, script := range scripts {
				loads, err := db.ParseInput(script)
				if err != nil {
					fmt.Println(cmd.Red("load all apps fail"))
					fmt.Println("    " + cmd.Red(err.Error()))
					return err
				}
				apps = append(apps, loads...)
			}
			if len(apps) == 0 {
				fmt.Println("No app found")
				return nil
			}
			for _, app := range apps {
				if app.ID == 0 {
					if err := db.SaveApp(app); err != nil {
						fmt.Println(cmd.Yellow(app.Name), cmd.Red("failed to saving"))
						fmt.Println("    " + cmd.Red(err.Error()))
						continue
					}
				}
				// todo options
				proc := subProcess(
					"-id="+strconv.Itoa(int(app.ID)),
					"-pmt=restart",
					"-sign="+xid.New().String(),
				)
				if err := proc.Wait(); err != nil {
					fmt.Println(cmd.Yellow(app.Name), cmd.Red("failed"))
					fmt.Println("    ", err.Error())
				} else {
					fmt.Println(cmd.Yellow(app.Name), cmd.Green("success"))
				}
			}
			fmt.Println("over")

			return nil
		}).

		// stop 命令
		Command("stop", "Stops a program").
		Argument("<names:...string>", "file/name/id").
		Action(func(c *cmd.Command, m map[string]any, a ...any) error {
			if len(a) != 1 {
				return c.ShowHelp()
			}
			var apps []*db.App
			scripts := a[0].([]string)
			for _, script := range scripts {
				loads, err := db.ParseInput(script)
				if err != nil {
					return err
				}
				apps = append(apps, loads...)
			}
			if len(apps) == 0 {
				fmt.Println("No app found")
				return nil
			}
			for _, app := range apps {
				if app.PID <= 0 {
					fmt.Println(cmd.Yellow(app.Name), "not running")
					return nil
				}
				proc := db.NewProc(app)
				err := proc.Stop()
				if err != nil {
					fmt.Println(cmd.Yellow(app.Name), cmd.Red("failed to stop"))
					fmt.Println("    " + cmd.Red(err.Error()))
					return err
				}
				fmt.Println(cmd.Yellow(app.Name), cmd.Green("stopped"))
				return nil
			}
			fmt.Println("over")
			return nil
		}).

		// 执行程序
		Run()
}

func subProcess(args ...string) *exec.Cmd {
	cmd := exec.Command(os.Args[0], args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Start()
	if err != nil {
		fmt.Fprintf(os.Stderr, "[-] Error: %s\n", err)
	}
	return cmd
}

// 移除不需要的参数
func strip(slice []string, remove string) []string {
	for i := 0; i < len(slice); {
		if slice[i] == remove {
			if i != len(slice)-1 {
				slice = append(slice[:i], slice[i+1:]...)
			} else {
				slice = slice[:i]
			}
		} else {
			i++
		}
	}
	return slice
}