Browse Source

tools/benchmark: ask for password when it is not supplied

harryge00 8 years ago
parent
commit
1c3567da90
2 changed files with 31 additions and 6 deletions
  1. 1 1
      tools/benchmark/cmd/root.go
  2. 30 5
      tools/benchmark/cmd/util.go

+ 1 - 1
tools/benchmark/cmd/root.go

@@ -67,7 +67,7 @@ func init() {
 	RootCmd.PersistentFlags().StringVar(&tls.KeyFile, "key", "", "identify HTTPS client using this SSL key file")
 	RootCmd.PersistentFlags().StringVar(&tls.CAFile, "cacert", "", "verify certificates of HTTPS-enabled servers using this CA bundle")
 
-	RootCmd.PersistentFlags().StringVar(&user, "user", "", "specify username and password in username:password format")
+	RootCmd.PersistentFlags().StringVar(&user, "user", "", "provide username[:password] and prompt if password is not supplied.")
 	RootCmd.PersistentFlags().DurationVar(&dialTimeout, "dial-timeout", 0, "dial timeout for client connections")
 
 	RootCmd.PersistentFlags().BoolVar(&targetLeader, "target-leader", false, "connect only to the leader node")

+ 30 - 5
tools/benchmark/cmd/util.go

@@ -24,6 +24,7 @@ import (
 	"github.com/coreos/etcd/clientv3"
 	"github.com/coreos/etcd/pkg/report"
 	"google.golang.org/grpc/grpclog"
+	"github.com/bgentry/speakeasy"
 )
 
 var (
@@ -33,6 +34,10 @@ var (
 
 	// leaderEps is a cache for holding endpoints of a leader node
 	leaderEps []string
+
+	// cache the username and password for multiple connections
+	globalUserName string
+	globalPassword string
 )
 
 func mustFindLeaderEndpoints(c *clientv3.Client) {
@@ -61,6 +66,26 @@ func mustFindLeaderEndpoints(c *clientv3.Client) {
 	os.Exit(1)
 }
 
+func getUsernamePassword(usernameFlag string) (string, string, error) {
+	if globalUserName != "" && globalPassword != "" {
+		return globalUserName, globalPassword, nil
+	}
+	colon := strings.Index(usernameFlag, ":")
+	if colon == -1 {
+		// Prompt for the password.
+		password, err := speakeasy.Ask("Password: ")
+		if err != nil {
+			return "", "", err
+		}
+		globalUserName = usernameFlag
+		globalPassword = password
+	} else {
+		globalUserName = usernameFlag[:colon]
+		globalPassword = usernameFlag[colon+1:]
+	}
+	return globalUserName, globalPassword, nil
+}
+
 func mustCreateConn() *clientv3.Client {
 	connEndpoints := leaderEps
 	if len(connEndpoints) == 0 {
@@ -81,14 +106,14 @@ func mustCreateConn() *clientv3.Client {
 	}
 
 	if len(user) != 0 {
-		splitted := strings.SplitN(user, ":", 2)
-		if len(splitted) != 2 {
-			fmt.Fprintf(os.Stderr, "bad user information: %s\n", user)
+		username, password, err := getUsernamePassword(user)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "bad user information: %s %v\n", user, err)
 			os.Exit(1)
 		}
+		cfg.Username = username
+		cfg.Password = password
 
-		cfg.Username = splitted[0]
-		cfg.Password = splitted[1]
 	}
 
 	client, err := clientv3.New(cfg)