feat:support execute command

cmd/ssx/cmd/root.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/spf13/cobra"
 
@@ -23,7 +24,22 @@ func NewRoot() *cobra.Command {
 		Use:   "ssx",
 		Short: "🦅 ssx is a retentive ssh client",
 		Example: `# If more than one flag of -i, -s ,-t specified, priority is ENTRY_ID > ADDRESS > TAG_NAME
-ssx [-i ENTRY_ID] [-s [USER@]HOST[:PORT]] [-k IDENTITY_FILE] [-t TAG_NAME]`,
+ssx [-i ENTRY_ID] [-s [USER@]HOST[:PORT]] [-k IDENTITY_FILE] [-t TAG_NAME]
+
+# You can also skip the parameters and log in directly with host or tag
+ssx [USER@]HOST[:PORT]
+ssx TAG_NAME
+
+# Fuzzy search is also supported
+# For example, you want to login to 192.168.1.100 and
+# suppose you can uniquely locate one entry by '100',
+# you just need to enter:
+ssx 100
+
+# If a command is specified, it will be executed on the remote host instead of a login shell.
+ssx 100 -c pwd
+# if the '-c' is omitted, the secend and subsequent arguments will be treated as COMMAND
+ssx 100 pwd`,
 		SilenceUsage:       true,
 		SilenceErrors:      true,
 		DisableAutoGenTag:  true,
@@ -49,6 +65,9 @@ ssx [-i ENTRY_ID] [-s [USER@]HOST[:PORT]] [-k IDENTITY_FILE] [-t TAG_NAME]`,
 				// just use first word as search key
 				opt.Keyword = args[0]
 			}
+			if len(args) > 1 && len(opt.Command) == 0 {
+				opt.Command = strings.Join(args[1:], " ")
+			}
 			return ssxInst.Main(cmd.Context())
 		},
 	}
@@ -57,6 +76,8 @@ ssx [-i ENTRY_ID] [-s [USER@]HOST[:PORT]] [-k IDENTITY_FILE] [-t TAG_NAME]`,
 	root.Flags().StringVarP(&opt.Addr, "server", "s", "", "target server address\nsupport formats: [user@]host[:port]")
 	root.Flags().StringVarP(&opt.Tag, "tag", "t", "", "search entry by tag")
 	root.Flags().StringVarP(&opt.IdentityFile, "keyfile", "k", "", "identity_file path")
+	root.Flags().StringVarP(&opt.Command, "cmd", "c", "", "the command to execute\nssh connection will exit after the execution complete")
+	root.Flags().DurationVar(&opt.Timeout, "timeout", 0, "timeout for connecting and executing command")
 
 	root.PersistentFlags().BoolVarP(&printVersion, "version", "v", false, "print ssx version")
 	root.PersistentFlags().BoolVar(&logVerbose, "verbose", false, "output detail logs")

ssx/client.go

@@ -10,9 +10,8 @@ import (
 	"sync"
 	"time"
 
-	"golang.org/x/crypto/ssh"
-
 	"github.com/containerd/console"
+	"golang.org/x/crypto/ssh"
 
 	"github.com/vimiix/ssx/internal/lg"
 	"github.com/vimiix/ssx/internal/terminal"
@@ -41,15 +40,44 @@ func (c *Client) touchEntry(e *entry.Entry) error {
 	return c.repo.TouchEntry(e)
 }
 
-func (c *Client) Run(ctx context.Context) error {
-	if err := c.login(ctx); err != nil {
+type ExecuteOption struct {
+	Command string
+	Stdout  io.Writer
+	Stderr  io.Writer
+	Timeout time.Duration
+}
+
+// Execute a command combined stdout and stderr output, then exit
+func (c *Client) Execute(ctx context.Context, opt *ExecuteOption) error {
+	if opt.Timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, opt.Timeout)
+		defer cancel()
+	}
+
+	if err := c.Login(ctx); err != nil {
 		return err
 	}
 	defer c.close()
 
-	if err := c.touchEntry(c.entry); err != nil {
+	sess, err := c.cli.NewSession()
+	if err != nil {
 		return err
 	}
+	defer sess.Close()
+
+	sess.Stdout = opt.Stdout
+	sess.Stderr = opt.Stderr
+	return sess.Run(opt.Command)
+}
+
+// Interact Bind the current terminal to provide an interactive interface
+func (c *Client) Interact(ctx context.Context) error {
+	if err := c.Login(ctx); err != nil {
+		return err
+	}
+	defer c.close()
+
 	lg.Info("connected server %s, version: %s",
 		c.entry.String(), string(c.cli.ServerVersion()))
 
@@ -151,22 +179,33 @@ func dialContext(ctx context.Context, network, addr string, config *ssh.ClientCo
 	return ssh.NewClient(c, chans, reqs), nil
 }
 
-func (c *Client) login(ctx context.Context) error {
+// Login connect remote server and touch enrty in storage
+func (c *Client) Login(ctx context.Context) error {
+	if err := c.connect(ctx); err != nil {
+		return err
+	}
+	if err := c.touchEntry(c.entry); err != nil {
+		lg.Error("failed to touch entry: %s", err)
+	}
+	return nil
+}
+
+func (c *Client) connect(ctx context.Context) error {
 	network := "tcp"
 	addr := net.JoinHostPort(c.entry.Host, c.entry.Port)
 	clientConfig, err := c.entry.GenSSHConfig(ctx)
 	if err != nil {
 		return err
 	}
-	lg.Info("connecting to %s", c.entry.String())
+	lg.Debug("connecting to %s", c.entry.String())
 	cli, err := dialContext(ctx, network, addr, clientConfig)
 	if err == nil {
 		c.cli = cli
 		return nil
 	}
 
 	if strings.Contains(err.Error(), "no supported methods remain") {
-		lg.Debug("failed login by default auth methods, try password again")
+		lg.Debug("failed connect by default auth methods, try password again")
 		fmt.Printf("%s@%s's password:", c.entry.User, c.entry.Host)
 		bs, readErr := terminal.ReadPassword(ctx)
 		fmt.Println()

ssx/ssx.go

@@ -11,6 +11,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/kevinburke/ssh_config"
 	"github.com/manifoldco/promptui"
@@ -33,6 +34,8 @@ type CmdOption struct {
 	Tag          string
 	IdentityFile string
 	Keyword      string
+	Command      string
+	Timeout      time.Duration
 }
 
 // Tidy complete unset fields with default values
@@ -158,7 +161,18 @@ func (s *SSX) Main(ctx context.Context) error {
 		return err
 	}
 
-	return NewClient(e, s.repo).Run(ctx)
+	client := NewClient(e, s.repo)
+	if len(s.opt.Command) > 0 {
+		opt := &ExecuteOption{
+			Command: s.opt.Command,
+			Stdout:  os.Stdout,
+			Stderr:  os.Stderr,
+			Timeout: s.opt.Timeout,
+		}
+		return client.Execute(ctx, opt)
+	}
+
+	return client.Interact(ctx)
 }
 
 func (s *SSX) getAllEntries() ([]*entry.Entry, error) {