From 75f2ecd8d1d349e5f5fbcf23465ba0fe7b9d76ef Mon Sep 17 00:00:00 2001
From: vimiix <i@vimiix.com>
Date: Thu, 14 Dec 2023 09:43:52 +0800
Subject: [PATCH] feat:support quick login without flags

---
 cmd/ssx/cmd/root.go |  5 +++++
 ssx/ssx.go          | 53 +++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/cmd/ssx/cmd/root.go b/cmd/ssx/cmd/root.go
index d18de1b..b0a360d 100644
--- a/cmd/ssx/cmd/root.go
+++ b/cmd/ssx/cmd/root.go
@@ -28,6 +28,7 @@ ssx [-i ENTRY_ID] [-s [USER@]HOST[:PORT]] [-k IDENTITY_FILE] [-t TAG_NAME]`,
 		SilenceErrors:      true,
 		DisableAutoGenTag:  true,
 		DisableSuggestions: true,
+		Args:               cobra.ArbitraryArgs, // accept arbitrary args for supporting quick login
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
 			lg.SetVerbose(logVerbose)
 			if !printVersion {
@@ -44,6 +45,10 @@ ssx [-i ENTRY_ID] [-s [USER@]HOST[:PORT]] [-k IDENTITY_FILE] [-t TAG_NAME]`,
 				fmt.Fprintln(os.Stdout, version.Detail())
 				return nil
 			}
+			if len(args) > 0 {
+				// just use first word as search key
+				opt.Keyword = args[0]
+			}
 			return ssxInst.Main(cmd.Context())
 		},
 	}
diff --git a/ssx/ssx.go b/ssx/ssx.go
index 3ef1bf4..9a2440d 100644
--- a/ssx/ssx.go
+++ b/ssx/ssx.go
@@ -32,6 +32,7 @@ type CmdOption struct {
 	Addr         string
 	Tag          string
 	IdentityFile string
+	Keyword      string
 }
 
 // Tidy complete unset fields with default values
@@ -142,7 +143,9 @@ func (s *SSX) Main(ctx context.Context) error {
 		e   *entry.Entry
 		err error
 	)
-	if s.opt.EntryID > 0 {
+	if s.opt.Keyword != "" {
+		e, err = s.searchEntry(s.opt.Keyword)
+	} else if s.opt.EntryID > 0 {
 		e, err = s.repo.GetEntry(s.opt.EntryID)
 	} else if s.opt.Addr != "" {
 		e, err = s.parseFuzzyAddr(s.opt.Addr)
@@ -158,7 +161,7 @@ func (s *SSX) Main(ctx context.Context) error {
 	return NewClient(e, s.repo).Run(ctx)
 }
 
-func (s *SSX) selectEntryFromAll() (*entry.Entry, error) {
+func (s *SSX) getAllEntries() ([]*entry.Entry, error) {
 	var es []*entry.Entry
 	em, err := s.repo.GetAllEntries()
 	if err != nil {
@@ -172,6 +175,52 @@ func (s *SSX) selectEntryFromAll() (*entry.Entry, error) {
 			es = append(es, e)
 		}
 	}
+	return es, nil
+}
+
+// search by host and tag first, if not found, then connect as a new entry
+func (s *SSX) searchEntry(keyword string) (*entry.Entry, error) {
+	es, err := s.getAllEntries()
+	if err != nil {
+		return nil, err
+	}
+	var candidates []*entry.Entry
+	for _, e := range es {
+		if strings.Contains(e.Host, keyword) ||
+			strings.Contains(strings.Join(e.Tags, " "), keyword) {
+			candidates = append(candidates, e)
+		}
+	}
+	if len(candidates) == 1 {
+		return candidates[0], nil
+	}
+	if len(candidates) > 1 {
+		return s.selectEntry(candidates)
+	}
+	lg.Debug("not found by keyword %q, treat it as new entry", keyword)
+	matches := addrRegex.FindStringSubmatch(keyword)
+	if len(matches) == 0 {
+		return nil, errors.Errorf("invalid address: %s", keyword)
+	}
+	username, host, port := matches[1], matches[2], matches[3]
+	e := &entry.Entry{
+		Host:    host,
+		User:    username,
+		Port:    port,
+		KeyPath: s.opt.IdentityFile,
+		Source:  entry.SourceSSXStore,
+	}
+	if err = e.Tidy(); err != nil {
+		return nil, err
+	}
+	return e, nil
+}
+
+func (s *SSX) selectEntryFromAll() (*entry.Entry, error) {
+	es, err := s.getAllEntries()
+	if err != nil {
+		return nil, err
+	}
 	return s.selectEntry(es)
 }