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 ", "Watches folder changes"). Option("-c, --cwd ", "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("", "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("", "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 }