use cobras MarkFlagRequired()

befehl.go

@@ -30,7 +30,7 @@ var store struct {
 	sshKey ssh.Signer
 }
 
-func Fire(targets *string, payload *string, routines *int, passedViper *viper.Viper) {
+func Fire(targets, payload string, routines int, passedViper *viper.Viper) {
 	Config = passedViper
 	bytePayload := readFile(payload)
 	populateSshKey()
@@ -47,7 +47,7 @@ func populateSshKey() {
 	if Config.GetString("auth.privatekeyfile") != "" {
 		privKeyFile = Config.GetString("auth.privatekeyfile")
 	}
-	rawKey := readFile(&privKeyFile)
+	rawKey := readFile(privKeyFile)
 	privKeyBytes, _ := pem.Decode(rawKey)
 
 	if x509.IsEncryptedPEMBlock(privKeyBytes) {
@@ -83,8 +83,8 @@ func (q *Queue) signifyComplete(total int) {
 	color.Magenta(fmt.Sprintf("Remaining: %d / %d\n", remaining, total))
 }
 
-func fireTorpedos(payload []byte, targets *string, routines *int) {
-	file, err := os.Open(*targets)
+func fireTorpedos(payload []byte, targets string, routines int) {
+	file, err := os.Open(targets)
 	if err != nil {
 		panic(err)
 	}
@@ -105,7 +105,7 @@ func fireTorpedos(payload []byte, targets *string, routines *int) {
 
 	var wg sync.WaitGroup
 	wg.Add(hostCnt)
-	var sem = make(chan int, *routines)
+	var sem = make(chan int, routines)
 
 	sshEntryUser := "root"
 	if Config.GetString("auth.sshuser") != "" {
@@ -217,8 +217,8 @@ func logPayloadRun(host string, output string) {
 	log.Printf("payload completed on %s! logfile at: %s\n", host, logFile)
 }
 
-func readFile(file *string) []byte {
-	read, err := ioutil.ReadFile(*file)
+func readFile(file string) []byte {
+	read, err := ioutil.ReadFile(file)
 	if err != nil {
 		panic(err)
 	}

cmd/execute.go

@@ -1,17 +1,12 @@
 package cmd
 
 import (
-	"fmt"
-
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 
 	"github.com/sgsullivan/befehl"
 )
 
-var payload string
-var routines int
-var hostsList string
-
 var executeCmd = &cobra.Command{
 	Use:   "execute",
 	Short: "Execute the given payload against the given hosts list",
@@ -38,22 +33,26 @@ sshuser = "eingeben"
 
 `,
 	Run: func(cmd *cobra.Command, args []string) {
-		for _, value := range []string{payload, hostsList} {
-			if value == "" {
-				panic("Missing payload, or hosts; see --help")
-			}
-			if routines == 0 {
-				fmt.Printf("--routines not given, defaulting to 30..\n")
-				routines = 30
-			}
+		hosts, _ := cmd.Flags().GetString("hosts")
+		payload, _ := cmd.Flags().GetString("payload")
+		routines, _ := cmd.Flags().GetInt("routines")
+
+		if routines == 0 {
+			color.Yellow("--routines not given, defaulting to 30..\n")
+			routines = 30
 		}
-		befehl.Fire(&hostsList, &payload, &routines, Config)
+
+		befehl.Fire(hosts, payload, routines, Config)
 	},
 }
 
 func init() {
 	RootCmd.AddCommand(executeCmd)
-	executeCmd.Flags().StringVarP(&payload, "payload", "", "", "file location to the payload, which contains the commands to execute on the remote hosts")
-	executeCmd.Flags().StringVarP(&hostsList, "hosts", "", "", "file location to hosts list, which contains all hosts (separated by newline) to run the payload on")
-	executeCmd.Flags().IntVarP(&routines, "routines", "", 0, "maximum number of payloads that will run at once (defaults to 30)")
+
+	executeCmd.Flags().String("payload", "", "file location to the payload, which contains the commands to execute on the remote hosts")
+	executeCmd.Flags().String("hosts", "", "file location to hosts list, which contains all hosts (separated by newline) to run the payload on")
+	executeCmd.Flags().Int("routines", 0, "maximum number of payloads that will run at once (defaults to 30)")
+
+	executeCmd.MarkFlagRequired("payload")
+	executeCmd.MarkFlagRequired("hosts")
 }

cmd/root.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
@@ -19,7 +20,9 @@ Dieses Werkzeug sollte mit Vorsicht verwendet werden; gegeben das Macht angebote
 
 func Execute() {
 	if err := RootCmd.Execute(); err != nil {
-		panic(err)
+		color.Red("A fatal error has occurred:")
+		fmt.Printf("%s\n", err)
+		os.Exit(1)
 	}
 }