Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
3 changed files
with
261 additions
and
85 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,130 @@ | ||
| package main | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "os" | ||
| "os/exec" | ||
| "strings" | ||
| "time" | ||
|
|
||
| "github.com/Sirupsen/logrus" | ||
| "github.com/justone/pmb/api" | ||
| ) | ||
|
|
||
| type RunCommand struct { | ||
| Message string `short:"m" long:"message" description:"Message to send."` | ||
| SendTrigger string `short:"s" long:"send-trigger" description:"Send trigger message when done."` | ||
| WaitTrigger string `short:"w" long:"wait-trigger" description:"Wait for trigger."` | ||
| TriggerAlways bool `short:"a" long:"trigger-always" description:"When trigger received, execute command if previous failed."` | ||
| Level float64 `short:"l" long:"level" description:"Notification level (1-5), higher numbers indictate higher importance" default:"3"` | ||
| } | ||
|
|
||
| var runCommand RunCommand | ||
|
|
||
| func (x *RunCommand) Execute(args []string) error { | ||
| bus := pmb.GetPMB(globalOptions.Primary) | ||
|
|
||
| if len(args) == 0 { | ||
| return fmt.Errorf("A command is required") | ||
| } | ||
|
|
||
| id := pmb.GenerateRandomID("run") | ||
|
|
||
| conn, err := bus.ConnectClient(id, !globalOptions.TrustKey) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| return runRun(conn, id, args) | ||
| } | ||
|
|
||
| func init() { | ||
| parser.AddCommand("run", | ||
| "Run a command.", | ||
| "", | ||
| &runCommand) | ||
| } | ||
|
|
||
| func runRun(conn *pmb.Connection, id string, args []string) error { | ||
|
|
||
| if waitTrigger := runCommand.WaitTrigger; len(waitTrigger) > 0 { | ||
| logrus.Infof("Waiting for trigger '%s' before starting...", waitTrigger) | ||
|
|
||
| var triggerMessage pmb.Message | ||
| WAIT: | ||
| for { | ||
| select { | ||
| case message := <-conn.In: | ||
| data := message.Contents | ||
| if data["type"].(string) == "Trigger" && data["from"].(string) == "run" && data["trigger"].(string) == waitTrigger { | ||
| triggerMessage = message | ||
| break WAIT | ||
| } | ||
| case _ = <-time.After(10 * time.Minute): | ||
| logrus.Warnf("Still waiting for trigger '%s'...", waitTrigger) | ||
| } | ||
| } | ||
|
|
||
| logrus.Infof("Trigger '%s' received...", waitTrigger) | ||
| note := pmb.Notification{ | ||
| Message: fmt.Sprintf("Received trigger %s", waitTrigger), | ||
| Level: 3, | ||
| } | ||
| pmb.SendNotification(conn, note) | ||
|
|
||
| if !runCommand.TriggerAlways && !triggerMessage.Contents["success"].(bool) { | ||
| return fmt.Errorf("Previous command failed, not running.") | ||
| } | ||
| } | ||
|
|
||
| message := runCommand.Message | ||
|
|
||
| cmd := exec.Command(args[0], args[1:]...) | ||
|
|
||
| cmd.Stdin = os.Stdin | ||
| cmd.Stdout = os.Stdout | ||
| cmd.Stderr = os.Stderr | ||
|
|
||
| command := strings.Join(args, " ") | ||
| logrus.Infof("Waiting for command '%s' to finish...", command) | ||
|
|
||
| err := cmd.Run() | ||
|
|
||
| cmdSuccess := true | ||
| result := "successfully" | ||
| if err != nil { | ||
| result = fmt.Sprintf("with error '%s'", err.Error()) | ||
| cmdSuccess = false | ||
| } | ||
| logrus.Infof("Process complete.") | ||
|
|
||
| if len(message) == 0 { | ||
| message = fmt.Sprintf("Command [%s] completed %s.", command, result) | ||
| } else { | ||
| message = fmt.Sprintf("%s. Command completed %s.", message, result) | ||
| } | ||
|
|
||
| note := pmb.Notification{Message: message, Level: runCommand.Level} | ||
| notifyErr := pmb.SendNotification(conn, note) | ||
|
|
||
| if sendTrigger := runCommand.SendTrigger; len(sendTrigger) > 0 { | ||
| logrus.Infof("Sending trigger '%s'.", sendTrigger) | ||
| note := pmb.Notification{ | ||
| Message: fmt.Sprintf("Sending trigger %s", sendTrigger), | ||
| Level: 3, | ||
| } | ||
| pmb.SendNotification(conn, note) | ||
|
|
||
| conn.Out <- pmb.Message{ | ||
| Contents: map[string]interface{}{ | ||
| "type": "Trigger", | ||
| "trigger": sendTrigger, | ||
| "from": "run", | ||
| "success": cmdSuccess, | ||
| }, | ||
| } | ||
| <-time.After(2 * time.Second) | ||
| } | ||
|
|
||
| return notifyErr | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,128 @@ | ||
| package main | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "os" | ||
| "os/exec" | ||
| "strconv" | ||
| "time" | ||
|
|
||
| "github.com/Sirupsen/logrus" | ||
| "github.com/justone/pmb/api" | ||
| ) | ||
|
|
||
| type WatchCommand struct { | ||
| Message string `short:"m" long:"message" description:"Message to send."` | ||
| Pid int `short:"p" long:"pid" description:"Notify after PID exits."` | ||
| File string `short:"f" long:"file" description:"Notify after file stops changing."` | ||
| Level float64 `short:"l" long:"level" description:"Notification level (1-5), higher numbers indictate higher importance" default:"3"` | ||
| } | ||
|
|
||
| var watchCommand WatchCommand | ||
|
|
||
| func (x *WatchCommand) Execute(args []string) error { | ||
| bus := pmb.GetPMB(globalOptions.Primary) | ||
|
|
||
| if len(watchCommand.Message) == 0 && watchCommand.Pid == 0 && watchCommand.File == "" { | ||
| return fmt.Errorf("A message or pid or file is required") | ||
| } | ||
|
|
||
| // fail fast if pid isn't found | ||
| if watchCommand.Pid > 0 { | ||
| found, _ := findProcess(watchCommand.Pid) | ||
|
|
||
| if !found { | ||
| return fmt.Errorf("Process %d not found.", watchCommand.Pid) | ||
| } | ||
| } | ||
| if len(watchCommand.File) > 0 { | ||
| if _, err := os.Stat(watchCommand.File); os.IsNotExist(err) { | ||
| return fmt.Errorf("File %s not found.", watchCommand.File) | ||
| } | ||
| } | ||
|
|
||
| id := pmb.GenerateRandomID("watch") | ||
|
|
||
| conn, err := bus.ConnectClient(id, !globalOptions.TrustKey) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| return runWatch(conn, id) | ||
| } | ||
|
|
||
| func init() { | ||
| parser.AddCommand("watch", | ||
| "Send a notification.", | ||
| "", | ||
| &watchCommand) | ||
| } | ||
|
|
||
| func runWatch(conn *pmb.Connection, id string) error { | ||
|
|
||
| message := watchCommand.Message | ||
|
|
||
| if watchCommand.Pid != 0 { | ||
|
|
||
| watchExecutable := "" | ||
| logrus.Infof("Waiting for pid %d to finish...\n", watchCommand.Pid) | ||
| for { | ||
| found, exec := findProcess(watchCommand.Pid) | ||
|
|
||
| // capture the name of the executable for the notification | ||
| if len(watchExecutable) == 0 { | ||
| watchExecutable = exec | ||
| } | ||
|
|
||
| if !found { | ||
| logrus.Infof("Process complete.") | ||
| break | ||
| } else { | ||
| time.Sleep(1 * time.Second) | ||
| } | ||
| } | ||
|
|
||
| if len(message) == 0 { | ||
| message = fmt.Sprintf("Command [%s] completed.", watchExecutable) | ||
| } | ||
| } | ||
| if len(watchCommand.File) > 0 { | ||
| var prevSize int64 | ||
| prevSize = -1 | ||
| for { | ||
| statInfo, _ := os.Stat(watchCommand.File) | ||
| if statInfo.Size() == prevSize { | ||
| logrus.Infof("File stabilized.") | ||
| break | ||
| } else { | ||
| prevSize = statInfo.Size() | ||
| time.Sleep(5 * time.Second) | ||
| } | ||
| } | ||
|
|
||
| if len(message) == 0 { | ||
| message = fmt.Sprintf("File [%s] stabilized.", watchCommand.File) | ||
| } | ||
| } | ||
|
|
||
| note := pmb.Notification{Message: message, Level: watchCommand.Level} | ||
| return pmb.SendNotification(conn, note) | ||
| } | ||
|
|
||
| // TODO: use a go-based library for this, maybe gopsutil | ||
| func findProcess(pid int) (bool, string) { | ||
|
|
||
| procCmd := exec.Command("/bin/ps", "-o", "pid=", "-p", strconv.Itoa(pid)) | ||
|
|
||
| err := procCmd.Run() | ||
|
|
||
| if _, ok := err.(*exec.ExitError); ok { | ||
| return false, "" | ||
| } else if err != nil { | ||
| return false, "" | ||
| } else { | ||
| return true, fmt.Sprintf("pid %d", pid) | ||
| } | ||
|
|
||
| return false, "" | ||
| } |