116 lines
2.7 KiB
Go
116 lines
2.7 KiB
Go
|
package cmd
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"strings"
|
||
|
|
"io"
|
||
|
|
"strconv"
|
||
|
|
"time"
|
||
|
|
"os"
|
||
|
|
"os/exec"
|
||
|
|
"os/signal"
|
||
|
|
"github.com/creack/pty"
|
||
|
|
"golang.org/x/term"
|
||
|
|
"syscall"
|
||
|
|
|
||
|
|
"github.com/spf13/cobra"
|
||
|
|
)
|
||
|
|
|
||
|
|
func init() {
|
||
|
|
rootCmd.AddCommand(timeCmd)
|
||
|
|
}
|
||
|
|
|
||
|
|
var timeCmd = &cobra.Command{
|
||
|
|
Use: "timer",
|
||
|
|
Short: "set the duration followed by the binary to be executed",
|
||
|
|
Long: "Set the timer duration after which the kill signal will be sent followed by the binary which should be executed.\n\nExamples:\nkilt timer 30 watch lscpu\nkilt timer 10 less /etc/passwd\nkilt timer 20 yes\nkilt timer 60 htop",
|
||
|
|
Args: cobra.MinimumNArgs(2), // Set at least <duration> and <binary>
|
||
|
|
Run: func(cmd *cobra.Command, args []string) {
|
||
|
|
if len(args) < 2 {
|
||
|
|
fmt.Println("Usage: timer <duration_in_seconds> <binary> [arguments...]")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
fmt.Println("Timer set to: " + strings.Join(args, " "))
|
||
|
|
if err:= runBinaryWithTimeout(args); err !=nil {
|
||
|
|
fmt.Println("Error:", err)
|
||
|
|
}
|
||
|
|
},
|
||
|
|
// Do not parse the flags which are meant for the started process, so that they are not interpreted as an arguemnt for kilt
|
||
|
|
DisableFlagParsing: true,
|
||
|
|
}
|
||
|
|
|
||
|
|
func resetTerminal() {
|
||
|
|
_ = exec.Command("stty", "sane").Run()
|
||
|
|
}
|
||
|
|
|
||
|
|
func runBinaryWithTimeout(args []string) error {
|
||
|
|
var durationStr string = args[0]
|
||
|
|
duration, err := strconv.Atoi(durationStr)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("Invalid duration: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
binaryPath := args[1]
|
||
|
|
binaryArgs := args[2:]
|
||
|
|
|
||
|
|
cmd := exec.Command(binaryPath, binaryArgs...)
|
||
|
|
|
||
|
|
oldState, err := term.MakeRaw(int(syscall.Stdin))
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("Failed to set terminal to raw mode: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
defer term.Restore(int(syscall.Stdin), oldState)
|
||
|
|
|
||
|
|
|
||
|
|
// cmd.Stdout = os.Stdout
|
||
|
|
// cmd.Stderr = os.Stderr
|
||
|
|
|
||
|
|
// if err := cmd.Start(); err != nil {
|
||
|
|
// return fmt.Errorf("Failed to start binary: %w", err)
|
||
|
|
// }
|
||
|
|
|
||
|
|
ptyFile, err := pty.Start(cmd)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("Failed to start pty: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
defer func() {
|
||
|
|
_ = ptyFile.Close()
|
||
|
|
term.Restore(int(syscall.Stdin), oldState)
|
||
|
|
resetTerminal()
|
||
|
|
}()
|
||
|
|
|
||
|
|
sigChan := make(chan os.Signal, 1)
|
||
|
|
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
||
|
|
go func() {
|
||
|
|
<- sigChan
|
||
|
|
_ = cmd.Process.Kill()
|
||
|
|
resetTerminal()
|
||
|
|
os.Exit(1)
|
||
|
|
}()
|
||
|
|
|
||
|
|
go func () {
|
||
|
|
_, _ = io.Copy(os.Stdout, ptyFile)
|
||
|
|
}()
|
||
|
|
|
||
|
|
done := make(chan error, 1)
|
||
|
|
go func() {
|
||
|
|
done <- cmd.Wait()
|
||
|
|
}()
|
||
|
|
|
||
|
|
select {
|
||
|
|
case <- time.After(time.Duration(duration) * time.Second):
|
||
|
|
if err := cmd.Process.Kill(); err != nil {
|
||
|
|
return fmt.Errorf("Failed to kill process: %w", err)
|
||
|
|
}
|
||
|
|
fmt.Println("Process killed as timeout has been reached.")
|
||
|
|
resetTerminal()
|
||
|
|
return nil
|
||
|
|
case err := <-done:
|
||
|
|
resetTerminal()
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|