From 76c9d8a5ae41f9964e31a4d7627c47656f2ae2ea Mon Sep 17 00:00:00 2001
From: Stargic <stargic@stargic.com>
Date: Sun, 12 Sep 2021 19:48:56 +0800
Subject: [PATCH] feat(climc): SSH Login with cloudroot (add new parameter
 --use-cloudroot in server-ssh)

fix(climc): Fixed the code logic error of FOR loop in server-ssh
---
 cmd/climc/shell/compute/servers.go      | 33 +++++++++++++++++++------
 pkg/mcclient/modules/mod_sshkeypairs.go |  5 ++++
 pkg/mcclient/options/servers.go         |  9 ++++---
 3 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/cmd/climc/shell/compute/servers.go b/cmd/climc/shell/compute/servers.go
index 4fd29caace..6757b9e49e 100644
--- a/cmd/climc/shell/compute/servers.go
+++ b/cmd/climc/shell/compute/servers.go
@@ -15,6 +15,7 @@
 package compute
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -818,13 +819,15 @@ func init() {
 			return fmt.Errorf("Not found ip address from server %s", opts.ID)
 		}
 
+		privateKey := ""
 		params := jsonutils.NewDict()
 		if len(opts.Key) > 0 {
-			privateKey, e := ioutil.ReadFile(opts.Key)
+			key, e := ioutil.ReadFile(opts.Key)
 			if e != nil {
 				return e
 			}
-			params.Add(jsonutils.NewString(string(privateKey)), "private_key")
+			params.Add(jsonutils.NewString(string(key)), "private_key")
+			privateKey = string(key)
 		}
 
 		i, e := modules.Servers.GetLoginInfo(s, srvid, params)
@@ -832,14 +835,14 @@ func init() {
 			return e
 		}
 		passwd, err := i.GetString("password")
-		if err != nil {
+		if err != nil && !opts.UseCloudroot {
 			return err
 		}
 		if opts.Password != "" {
 			passwd = opts.Password
 		}
 		user, err := i.GetString("username")
-		if err != nil {
+		if err != nil && !opts.UseCloudroot {
 			return err
 		}
 		if opts.User != "" {
@@ -869,8 +872,22 @@ func init() {
 			}
 		}
 
+		if opts.UseCloudroot {
+			var err error
+			privateKey, err = modules.Sshkeypairs.FetchPrivateKeyBySession(context.Background(), s)
+			if err != nil {
+				return err
+			}
+			passwd = ""
+			user = "cloudroot"
+		}
+
 		var sshCli *ssh.Client
-		for sshCli, err = ssh.NewClient(host, port, user, passwd, ""); err != nil; {
+		err = nil
+		for ; sshCli == nil; sshCli, err = ssh.NewClient(host, port, user, passwd, privateKey) {
+			if err == nil {
+				continue
+			}
 			if opts.Host != "" {
 				return err
 			}
@@ -879,9 +896,9 @@ func init() {
 				return err
 			} else {
 				if vpcid != "default" {
-					forwardItem, err = openForward(s, srvid)
-					if err != nil {
-						return err
+					forwardItem, e = openForward(s, srvid)
+					if e != nil {
+						return e
 					}
 					host = forwardItem.ProxyAddr
 					port = forwardItem.ProxyPort
diff --git a/pkg/mcclient/modules/mod_sshkeypairs.go b/pkg/mcclient/modules/mod_sshkeypairs.go
index b2c8b4d6dc..fc7e84a8af 100644
--- a/pkg/mcclient/modules/mod_sshkeypairs.go
+++ b/pkg/mcclient/modules/mod_sshkeypairs.go
@@ -48,6 +48,11 @@ func (this *SSshkeypairManager) List(s *mcclient.ClientSession, params jsonutils
 
 func (this *SSshkeypairManager) FetchPrivateKey(ctx context.Context, userCred mcclient.TokenCredential) (string, error) {
 	s := auth.GetSession(ctx, userCred, "", "")
+	return this.FetchPrivateKeyBySession(ctx, s)
+}
+
+func (this *SSshkeypairManager) FetchPrivateKeyBySession(ctx context.Context, s *mcclient.ClientSession) (string, error) {
+	userCred := s.GetToken()
 	jd := jsonutils.NewDict()
 	var jr jsonutils.JSONObject
 	if userCred.HasSystemAdminPrivilege() {
diff --git a/pkg/mcclient/options/servers.go b/pkg/mcclient/options/servers.go
index d3c4b6a19f..1c8ab2077b 100644
--- a/pkg/mcclient/options/servers.go
+++ b/pkg/mcclient/options/servers.go
@@ -99,10 +99,11 @@ type ServerLoginInfoOptions struct {
 
 type ServerSSHLoginOptions struct {
 	ServerLoginInfoOptions
-	Host     string `help:"IP address or hostname of the server"`
-	Port     int    `help:"SSH service port" default:"22"`
-	User     string `help:"SSH login user"`
-	Password string `help:"SSH password"`
+	Host         string `help:"IP address or hostname of the server"`
+	Port         int    `help:"SSH service port" default:"22"`
+	User         string `help:"SSH login user"`
+	Password     string `help:"SSH password"`
+	UseCloudroot bool   `help:"SSH login with cloudroot"`
 }
 
 type ServerConvertToKvmOptions struct {