Build functionality to connect to host and run task (via 'connectToHo…

…st' function).

cmd/args.go

@@ -30,7 +30,7 @@ var ArgsCmd = &cobra.Command{
 		handleErr(err)
 
 		// Get host.
-		proxy, err := konnect.Get(host)
+		proxy, err := konnect.GetHost(host)
 		handleErr(err)
 
 		// Get args for host.

cmd/cmd.go

@@ -90,7 +90,7 @@ func InteractivePrompt(cmd *cobra.Command) {
 	}
 
 	// Get proxy.
-	proxy, err := konnect.Get(answer.Hostname)
+	proxy, err := konnect.GetHost(answer.Hostname)
 	handleErr(err)
 
 	// Connect to host.
@@ -101,6 +101,7 @@ func InteractivePrompt(cmd *cobra.Command) {
 
 // AddCommands - Connects subcommands to the RootCmd.
 func AddCommands() {
+	ConnectCmd.AddCommand(TaskCmd)
 	RootCmd.AddCommand(ArgsCmd)
 	RootCmd.AddCommand(ConnectCmd)
 	RootCmd.AddCommand(InitCmd)

cmd/connect.go

@@ -1,39 +1,69 @@
 package cmd
 
 import (
+	"fmt"
 	"log"
+	"os"
 
-	"github.com/exitshell/konnect/engine"
 	"github.com/spf13/cobra"
 )
 
+var hostName string
+var taskName string
+
 // ConnectCmd - Connect to a host.
 var ConnectCmd = &cobra.Command{
 	Use:   "to",
 	Short: "Connect to a host",
 	Long:  "Connect to a host",
 	Run: func(cmd *cobra.Command, args []string) {
-		// Resolve filename from flags.
-		filename, err := resolveFilename(cmd)
-		handleErr(err)
+		fmt.Printf("[connectcmd] %v (%T)\n", args, args)
 
-		// Check that only one host was specified.
-		if len(args) != 1 {
-			log.Fatal("Please specify one host")
+		if len(args) == 0 {
+			log.Fatal("Please specify a host")
 		}
-		host := args[0]
 
-		// Init engine.
-		konnect, err := engine.Init(filename)
-		handleErr(err)
+		// Set hostname.
+		hostName = args[0]
 
-		// Get host.
-		proxy, err := konnect.Get(host)
+		if len(args) == 1 {
+			// Connect to host.
+			err := connectToHost(cmd, hostName, "")
+			handleErr(err)
+			os.Exit(0)
+		}
+
+		// Find the subcommand.
+		subCmd, subArgs, err := cmd.Find(args[1:])
 		handleErr(err)
+		// If the subcommand is the same as the original command,
+		// then no new subcommand was found. In that case, exit
+		// with an error.
+		if subCmd.Use == cmd.Use {
+			log.Fatalf("Cannot parse subcommand %v\n", subArgs)
+		} else {
+			subCmd.Run(cmd, subArgs)
+		}
+	},
+}
+
+// TaskCmd - Run a task on a host.
+var TaskCmd = &cobra.Command{
+	Use:   "and",
+	Short: "Run a task on a host",
+	Long:  "Run a task on a host",
+	Run: func(cmd *cobra.Command, args []string) {
+		fmt.Printf("[taskcmd] %v (%T)\n", args, args)
 
-		// Connect to host.
-		if err := proxy.Connect(); err != nil {
-			log.Fatal(err)
+		if len(args) == 0 {
+			log.Fatal("Please specify a task")
 		}
+
+		// Set task name.
+		taskName = args[0]
+
+		// Connect to host and run a command.
+		err := connectToHost(cmd, hostName, taskName)
+		handleErr(err)
 	},
 }

cmd/utils_.go

@@ -123,3 +123,33 @@ func makeDefaultConfig(filename string) error {
 
 	return nil
 }
+
+func connectToHost(cmd *cobra.Command, hostName, taskName string) error {
+	// Resolve filename from flags.
+	filename, err := resolveFilename(cmd)
+	handleErr(err)
+
+	// Init engine.
+	konnect, err := engine.Init(filename)
+	handleErr(err)
+
+	// Get host.
+	proxy, err := konnect.GetHost(hostName)
+	handleErr(err)
+
+	// Get task if specified.
+	if taskName != "" {
+		task, err := konnect.GetTask(taskName)
+		handleErr(err)
+
+		// Add task command to proxy.
+		proxy.ExtraArgs = task.Command
+	}
+
+	// Connect to host.
+	if err := proxy.Connect(); err != nil {
+		log.Fatal(err)
+	}
+
+	return nil
+}

engine/engine.go

@@ -22,16 +22,27 @@ type Konnect struct {
 }
 
 // Get an SSHProxy object by name.
-func (k *Konnect) Get(name string) (*proxy.SSHProxy, error) {
+func (k *Konnect) GetHost(name string) (*proxy.SSHProxy, error) {
 	proxy, ok := k.Hosts[name]
 	// Return error if SSHProxy rule is not found.
 	if !ok {
-		errMsg := fmt.Sprintf("Host '%v' not found", name)
+		errMsg := fmt.Sprintf("Host '%v' not found\n", name)
 		return proxy, errors.New(errMsg)
 	}
 	return proxy, nil
 }
 
+// GetTask - Get an SSHTask object by name.
+func (k *Konnect) GetTask(name string) (*task.SSHTask, error) {
+	task, ok := k.Tasks[name]
+	// Return error if SSHTask is not found.
+	if !ok {
+		errMsg := fmt.Sprintf("Task '%v' not found\n", name)
+		return task, errors.New(errMsg)
+	}
+	return task, nil
+}
+
 // GetHostNames - Get host names in sorted order (asc).
 func (k *Konnect) GetHostNames() []string {
 	names := []string{}
@@ -101,9 +112,9 @@ func (k *Konnect) UnmarshalHosts(byteStr []byte) error {
 	for key, val := range tempHosts.Hosts {
 		switch val.(type) {
 
-		// If the host value is a string, then it means that an
+		// If the value is a string, then it means that an
 		// SSHProxy.Host value was supplied only. In this case,
-		// we create an SSHProxy with thie value as the `Host`.
+		// we create an SSHProxy with this value as the `Host`.
 		case string:
 			// Construct an SSHProxy object.
 			proxy := proxy.New("", val.(string), 0, "")
@@ -112,7 +123,7 @@ func (k *Konnect) UnmarshalHosts(byteStr []byte) error {
 			// Assign to Konnect.
 			k.Hosts[key] = proxy
 
-		// If the host value is an interfact map, then it means
+		// If the value is an interfact map, then it means
 		// that an SSHProxy was possibly defined in full. In
 		// this case, we marshal the map to a byte string, and
 		// unmarhsal the byte string into an SSHProxy object.

proxy/proxy.go

@@ -29,6 +29,8 @@ type SSHProxy struct {
 	Name string `yaml:"-"`
 	// A bool to determine if the connection is ok.
 	Connection bool `yaml:"-"`
+	// Extra args.
+	ExtraArgs string `yaml:"-"`
 }
 
 // String representation of an SSHProxy object.
@@ -61,21 +63,28 @@ func (s *SSHProxy) PrintStatus() string {
 	return fmt.Sprintf("Connection %v\t-> [%v]", status(connectionStr), s.Name)
 }
 
-// Args - Return the full SSH command for SSHProxy.
+// Args - Return the full SSH command as a string slice of args.
 func (s *SSHProxy) Args() []string {
-	return []string{
+	args := []string{
 		"ssh",
 		"-i",
 		s.Key,
 		"-p",
 		strconv.Itoa(s.Port),
 		fmt.Sprintf("%v@%v", s.User, s.Host),
 	}
+	// If SSHProxy has extra args, then
+	// add them to the `args` slice.
+	if len(s.ExtraArgs) > 0 {
+		args = append(args, strings.Fields(s.ExtraArgs)...)
+	}
+	return args
 }
 
 // Connect to host.
 func (s *SSHProxy) Connect() error {
 	args := s.Args()
+	fmt.Printf("\nWould run: %v %v\n", args[0], strings.Join(args[1:], " "))
 
 	// Make command, and pipe streams.
 	cmd := exec.Command(args[0], args[1:]...)
@@ -85,6 +94,7 @@ func (s *SSHProxy) Connect() error {
 
 	// Run command.
 	return cmd.Run()
+	// return nil
 }
 
 // Validate SSHProxy fields.

task/task.go

@@ -1,11 +1,14 @@
 package task
 
-import "fmt"
+import (
+	"fmt"
+	"strings"
+)
 
-// type SSHTask string
+// SSHTask -
 type SSHTask struct {
 	Command string
-	Name string
+	Name    string
 }
 
 // String representation of an SSHTask object.
@@ -17,5 +20,10 @@ func (t *SSHTask) String() string {
 func (t *SSHTask) Info() string {
 	return fmt.Sprintf("[%v]\n"+
 		"  Command: %v\n",
-	t.Name, t.Command)
+		t.Name, t.Command)
+}
+
+// Args - Return the SSHTask command as a string slice.
+func (t *SSHTask) Args() []string {
+	return strings.Fields(t.Command)
 }