Add: Certificate verification login function

Driver-C · Driver-C · commit affa7088995b · 2024-07-10T16:09:50.000+08:00
diff --git a/cmd/create/create.go b/cmd/create/create.go
@@ -19,5 +19,6 @@ func NewCreateCommand() *cobra.Command {
 	createCmd.AddCommand(ports.NewPortsCommand())
 	createCmd.AddCommand(passwords.NewPasswordsCommand())
 	createCmd.AddCommand(caches.NewCachesCommand())
+	createCmd.AddCommand(NewKeysCommand())
 	return createCmd
 }
diff --git a/cmd/create/keys.go b/cmd/create/keys.go
@@ -0,0 +1,26 @@
+package create
+
+import (
+	"github.com/Driver-C/tryssh/config"
+	"github.com/Driver-C/tryssh/control/create"
+	"github.com/spf13/cobra"
+)
+
+const createType = "keys"
+
+func NewKeysCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "keys <keyFilePath>",
+		Args:    cobra.ExactArgs(1),
+		Short:   "Create a alternate key file path",
+		Long:    "Create a alternate key file path",
+		Aliases: []string{"key"},
+		Run: func(cmd *cobra.Command, args []string) {
+			keyPath := args[0]
+			configuration := config.LoadConfig()
+			ctl := create.NewCreateController(createType, keyPath, configuration)
+			ctl.ExecuteCreate()
+		},
+	}
+	return cmd
+}
diff --git a/cmd/delete/delete.go b/cmd/delete/delete.go
@@ -19,5 +19,6 @@ func NewDeleteCommand() *cobra.Command {
 	deleteCmd.AddCommand(ports.NewPortsCommand())
 	deleteCmd.AddCommand(passwords.NewPasswordsCommand())
 	deleteCmd.AddCommand(caches.NewCachesCommand())
+	deleteCmd.AddCommand(NewKeysCommand())
 	return deleteCmd
 }
diff --git a/cmd/delete/keys.go b/cmd/delete/keys.go
@@ -0,0 +1,26 @@
+package delete
+
+import (
+	"github.com/Driver-C/tryssh/config"
+	"github.com/Driver-C/tryssh/control/delete"
+	"github.com/spf13/cobra"
+)
+
+const deleteType = "keys"
+
+func NewKeysCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "keys <keyFilePath>",
+		Args:    cobra.ExactArgs(1),
+		Short:   "Delete a alternate key file path",
+		Long:    "Delete a alternate key file path",
+		Aliases: []string{"key"},
+		Run: func(cmd *cobra.Command, args []string) {
+			keyPath := args[0]
+			configuration := config.LoadConfig()
+			ctl := delete.NewDeleteController(deleteType, keyPath, configuration)
+			ctl.ExecuteDelete()
+		},
+	}
+	return cmd
+}
diff --git a/cmd/get/get.go b/cmd/get/get.go
@@ -18,5 +18,6 @@ func NewGetCommand() *cobra.Command {
 	getCmd.AddCommand(ports.NewPortsCommand())
 	getCmd.AddCommand(passwords.NewPasswordsCommand())
 	getCmd.AddCommand(caches.NewCachesCommand())
+	getCmd.AddCommand(NewKeysCommand())
 	return getCmd
 }
diff --git a/cmd/get/keys.go b/cmd/get/keys.go
@@ -0,0 +1,28 @@
+package get
+
+import (
+	"github.com/Driver-C/tryssh/config"
+	"github.com/Driver-C/tryssh/control/get"
+	"github.com/spf13/cobra"
+)
+
+const getType = "keys"
+
+func NewKeysCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "keys <keyFilePath>",
+		Short:   "Get alternate key file path",
+		Long:    "Get alternate key file path",
+		Aliases: []string{"key"},
+		Run: func(cmd *cobra.Command, args []string) {
+			var keyPath string
+			if len(args) > 0 {
+				keyPath = args[0]
+			}
+			configuration := config.LoadConfig()
+			ctl := get.NewGetController(getType, keyPath, configuration)
+			ctl.ExecuteGet()
+		},
+	}
+	return cmd
+}
diff --git a/config/config.go b/config/config.go
@@ -38,6 +38,7 @@ type MainConfig struct {
 		Ports     []string `yaml:"ports,flow"`
 		Users     []string `yaml:"users,flow"`
 		Passwords []string `yaml:"passwords,flow"`
+		Keys      []string `yaml:"keys,flow"`
 	} `yaml:"main"`
 	ServerLists []ServerListConfig `yaml:"serverList"`
 }
@@ -48,6 +49,7 @@ type ServerListConfig struct {
 	Port     string `yaml:"port"`
 	User     string `yaml:"user"`
 	Password string `yaml:"password"`
+	Key      string `yaml:"key"`
 	Alias    string `yaml:"alias"`
 }
 
@@ -122,7 +124,8 @@ func GenerateCombination(ip string, user string, conf *MainConfig) (combinations
 		users = utils.InterfaceSlice(conf.Main.Users)
 	}
 	passwords := utils.InterfaceSlice(conf.Main.Passwords)
+	keys := utils.InterfaceSlice(conf.Main.Keys)
 	// Generate combinations with immutable parameter order
-	combinations = cartesian.Iter(ips, ports, users, passwords)
+	combinations = cartesian.Iter(ips, ports, users, passwords, keys)
 	return
 }
diff --git a/control/create/create.go b/control/create/create.go
@@ -11,6 +11,7 @@ const (
 	typePorts     = "ports"
 	typePasswords = "passwords"
 	typeCaches    = "caches"
+	typeKeys      = "keys"
 )
 
 type CacheContent struct {
@@ -41,6 +42,10 @@ func (cc Controller) ExecuteCreate() {
 		cc.configuration.Main.Passwords = utils.RemoveDuplicate(
 			append(cc.configuration.Main.Passwords, cc.createContent))
 		cc.updateConfig()
+	case typeKeys:
+		cc.configuration.Main.Keys = utils.RemoveDuplicate(
+			append(cc.configuration.Main.Keys, cc.createContent))
+		cc.updateConfig()
 	case typeCaches:
 		cc.createCaches()
 		cc.updateConfig()
diff --git a/control/delete/delete.go b/control/delete/delete.go
@@ -10,6 +10,7 @@ const (
 	typePorts     = "ports"
 	typePasswords = "passwords"
 	typeCaches    = "caches"
+	typeKey       = "keys"
 )
 
 type Controller struct {
@@ -44,6 +45,14 @@ func (dc Controller) ExecuteDelete() {
 		} else {
 			utils.Logger.Warnf("No matching password: %s\n", dc.deleteContent)
 		}
+	case typeKey:
+		contents := dc.configuration.Main.Keys
+		if newContents := dc.searchAndDelete(contents); newContents != nil {
+			dc.configuration.Main.Keys = newContents
+			dc.updateConfig()
+		} else {
+			utils.Logger.Warnf("No matching key: %s\n", dc.deleteContent)
+		}
 	case typeCaches:
 		// dc.deleteContent is ipAddress
 		var deleteCount int
diff --git a/control/get/get.go b/control/get/get.go
@@ -10,6 +10,7 @@ const (
 	typePorts     = "ports"
 	typePasswords = "passwords"
 	typeCaches    = "caches"
+	typeKeys      = "keys"
 )
 
 type Controller struct {
@@ -29,6 +30,9 @@ func (gc Controller) ExecuteGet() {
 	case typePasswords:
 		fmt.Println("INDEX	PASSWORD")
 		gc.searchAndPrint(gc.configuration.Main.Passwords)
+	case typeKeys:
+		fmt.Println("INDEX	KEY")
+		gc.searchAndPrint(gc.configuration.Main.Keys)
 	case typeCaches:
 		// gc.getContent is ipAddress
 		fmt.Println("INDEX	CACHE")
diff --git a/launcher/base.go b/launcher/base.go
@@ -18,6 +18,8 @@ const (
 	SSHKeyKeyword        = "SSH-KEY"
 )
 
+var keysMap = sync.Map{}
+
 type Connector interface {
 	Launch() bool
 	CreateConnection() (sshClient *ssh.Client, err error)
@@ -30,6 +32,7 @@ type SshConnector struct {
 	Port         string
 	User         string
 	Password     string
+	Key          string
 	SshTimeout   time.Duration
 	HostKeyMutex *sync.Mutex
 }
@@ -43,11 +46,30 @@ func (sc *SshConnector) LoadConfig() (config *ssh.ClientConfig) {
 	if sc.HostKeyMutex == nil {
 		sc.HostKeyMutex = new(sync.Mutex)
 	}
+
+	var authMethods []ssh.AuthMethod
+	var privateKey []byte
+	if sc.Key != "" {
+		if _, ok := keysMap.Load(sc.Key); !ok {
+			if pk, status := utils.ReadFile(sc.Key); status {
+				keysMap.Store(sc.Key, pk)
+				privateKey = pk
+			}
+		} else {
+			pk, _ := keysMap.Load(sc.Key)
+			privateKey = pk.([]byte)
+		}
+		signer, err := ssh.ParsePrivateKey(privateKey)
+		if err == nil {
+			authMethods = append(authMethods, ssh.PublicKeys(signer))
+		} else {
+			utils.Logger.Errorln("Failed to parse private key: %v", err)
+		}
+	}
+	authMethods = append(authMethods, ssh.Password(sc.Password))
 	config = &ssh.ClientConfig{
-		User: sc.User,
-		Auth: []ssh.AuthMethod{
-			ssh.Password(sc.Password),
-		},
+		User:            sc.User,
+		Auth:            authMethods,
 		HostKeyCallback: trustedHostKeyCallback(searchKeyFromAddress(sc.Ip), sc.Ip, sc.HostKeyMutex),
 		Timeout:         sc.SshTimeout,
 	}
@@ -95,6 +117,7 @@ func GetSshConnectorFromConfig(conf *config.ServerListConfig) *SshConnector {
 		Port:     conf.Port,
 		User:     conf.User,
 		Password: conf.Password,
+		Key:      conf.Key,
 	}
 }
 
@@ -105,6 +128,7 @@ func GetConfigFromSshConnector(tgt *SshConnector) *config.ServerListConfig {
 		Port:     tgt.Port,
 		User:     tgt.User,
 		Password: tgt.Password,
+		Key:      tgt.Key,
 	}
 }
 
diff --git a/launcher/scp/scp.go b/launcher/scp/scp.go
@@ -62,6 +62,7 @@ func NewScpLaunchersByCombinations(combinations chan []interface{}, src string,
 				Port:       com[1].(string),
 				User:       com[2].(string),
 				Password:   com[3].(string),
+				Key:        com[4].(string),
 				SshTimeout: sshTimeout,
 			},
 			Src:       src,
diff --git a/launcher/ssh/ssh.go b/launcher/ssh/ssh.go
@@ -25,6 +25,7 @@ func NewSshLaunchersByCombinations(combinations chan []interface{},
 			Port:       com[1].(string),
 			User:       com[2].(string),
 			Password:   com[3].(string),
+			Key:        com[4].(string),
 			SshTimeout: sshTimeout,
 		}})
 	}

Original file line number	Diff line number	Diff line change
`@@ -19,5 +19,6 @@ func NewCreateCommand() *cobra.Command {`
`19`	`19`	`createCmd.AddCommand(ports.NewPortsCommand())`
`20`	`20`	`createCmd.AddCommand(passwords.NewPasswordsCommand())`
`21`	`21`	`createCmd.AddCommand(caches.NewCachesCommand())`
	`22`	`+ createCmd.AddCommand(NewKeysCommand())`
`22`	`23`	`return createCmd`
`23`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -19,5 +19,6 @@ func NewDeleteCommand() *cobra.Command {`
`19`	`19`	`deleteCmd.AddCommand(ports.NewPortsCommand())`
`20`	`20`	`deleteCmd.AddCommand(passwords.NewPasswordsCommand())`
`21`	`21`	`deleteCmd.AddCommand(caches.NewCachesCommand())`
	`22`	`+ deleteCmd.AddCommand(NewKeysCommand())`
`22`	`23`	`return deleteCmd`
`23`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,5 +18,6 @@ func NewGetCommand() *cobra.Command {`
`18`	`18`	`getCmd.AddCommand(ports.NewPortsCommand())`
`19`	`19`	`getCmd.AddCommand(passwords.NewPasswordsCommand())`
`20`	`20`	`getCmd.AddCommand(caches.NewCachesCommand())`
	`21`	`+ getCmd.AddCommand(NewKeysCommand())`
`21`	`22`	`return getCmd`
`22`	`23`	`}`
Original file line number	Diff line number	Diff line change
`@@ -38,6 +38,7 @@ type MainConfig struct {`
`38`	`38`	Ports []string `yaml:"ports,flow"`
`39`	`39`	Users []string `yaml:"users,flow"`
`40`	`40`	Passwords []string `yaml:"passwords,flow"`
	`41`	+ Keys []string `yaml:"keys,flow"`
`41`	`42`	} `yaml:"main"`
`42`	`43`	ServerLists []ServerListConfig `yaml:"serverList"`
`43`	`44`	`}`
`@@ -48,6 +49,7 @@ type ServerListConfig struct {`
`48`	`49`	Port string `yaml:"port"`
`49`	`50`	User string `yaml:"user"`
`50`	`51`	Password string `yaml:"password"`
	`52`	+ Key string `yaml:"key"`
`51`	`53`	Alias string `yaml:"alias"`
`52`	`54`	`}`
`53`	`55`
`@@ -122,7 +124,8 @@ func GenerateCombination(ip string, user string, conf *MainConfig) (combinations`
`122`	`124`	`users = utils.InterfaceSlice(conf.Main.Users)`
`123`	`125`	`}`
`124`	`126`	`passwords := utils.InterfaceSlice(conf.Main.Passwords)`
	`127`	`+ keys := utils.InterfaceSlice(conf.Main.Keys)`
`125`	`128`	`// Generate combinations with immutable parameter order`
`126`		`- combinations = cartesian.Iter(ips, ports, users, passwords)`
	`129`	`+ combinations = cartesian.Iter(ips, ports, users, passwords, keys)`
`127`	`130`	`return`
`128`	`131`	`}`