login: allow 2fa code input if mandated (#175)

Signed-off-by: Abhishek Kumar <[email protected]>

Author: shwstppr

cmd/network.go

@@ -49,6 +49,84 @@ func findSessionCookie(cookies []*http.Cookie) *http.Cookie {
 	return nil
 }
 
+func getLoginResponse(responseBody []byte) (map[string]interface{}, error) {
+	var responseMap map[string]interface{}
+	err := json.Unmarshal(responseBody, &responseMap)
+	if err != nil {
+		return nil, errors.New("failed to parse login response: " + err.Error())
+	}
+	loginRespRaw, ok := responseMap["loginresponse"]
+	if !ok {
+		return nil, errors.New("failed to parse login response, expected 'loginresponse' key not found")
+	}
+	loginResponse, ok := loginRespRaw.(map[string]interface{})
+	if !ok {
+		return nil, errors.New("failed to parse login response, expected 'loginresponse' to be a map")
+	}
+	return loginResponse, nil
+}
+
+func getResponseBooleanValue(response map[string]interface{}, key string) (bool, bool) {
+	v, found := response[key]
+	if !found {
+		return false, false
+	}
+	switch value := v.(type) {
+	case bool:
+		return true, value
+	case string:
+		return true, strings.ToLower(value) == "true"
+	case float64:
+		return true, value != 0
+	default:
+		return true, false
+	}
+}
+
+func checkLogin2FAPromptAndValidate(r *Request, response map[string]interface{}, sessionKey string) error {
+	if !r.Config.HasShell {
+		return nil
+	}
+	config.Debug("Checking if 2FA is enabled and verified for the user ", response)
+	found, is2faEnabled := getResponseBooleanValue(response, "is2faenabled")
+	if !found || !is2faEnabled {
+		config.Debug("2FA is not enabled for the user, skipping 2FA validation")
+		return nil
+	}
+	found, is2faVerified := getResponseBooleanValue(response, "is2faverified")
+	if !found || is2faVerified {
+		config.Debug("2FA is already verified for the user, skipping 2FA validation")
+		return nil
+	}
+	activeSpinners := r.Config.PauseActiveSpinners()
+	fmt.Print("Enter 2FA code: ")
+	var code string
+	fmt.Scanln(&code)
+	if activeSpinners > 0 {
+		r.Config.ResumePausedSpinners()
+	}
+	params := make(url.Values)
+	params.Add("command", "validateUserTwoFactorAuthenticationCode")
+	params.Add("codefor2fa", code)
+	params.Add("sessionkey", sessionKey)
+
+	msURL, _ := url.Parse(r.Config.ActiveProfile.URL)
+
+	config.Debug("Validating 2FA with POST URL:", msURL, params)
+	spinner := r.Config.StartSpinner("trying to validate 2FA...")
+	resp, err := r.Client().PostForm(msURL.String(), params)
+	r.Config.StopSpinner(spinner)
+	if err != nil {
+		return errors.New("failed to failed to validate 2FA code: " + err.Error())
+	}
+	config.Debug("ValidateUserTwoFactorAuthenticationCode POST response status code:", resp.StatusCode)
+	if resp.StatusCode != http.StatusOK {
+		r.Client().Jar, _ = cookiejar.New(nil)
+		return errors.New("failed to validate 2FA code, please check the code. Invalidating session")
+	}
+	return nil
+}
+
 // Login logs in a user based on provided request and returns http client and session key
 func Login(r *Request) (string, error) {
 	params := make(url.Values)
@@ -81,6 +159,13 @@ func Login(r *Request) (string, error) {
 		return "", e
 	}
 
+	body, _ := ioutil.ReadAll(resp.Body)
+	config.Debug("Login response body:", string(body))
+	loginResponse, err := getLoginResponse(body)
+	if err != nil {
+		return "", err
+	}
+
 	var sessionKey string
 	curTime := time.Now()
 	expiryDuration := 15 * time.Minute
@@ -98,6 +183,9 @@ func Login(r *Request) (string, error) {
 	}()
 
 	config.Debug("Login sessionkey:", sessionKey)
+	if err := checkLogin2FAPromptAndValidate(r, loginResponse, sessionKey); err != nil {
+		return "", err
+	}
 	return sessionKey, nil
 }
 

config/config.go

@@ -30,6 +30,7 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/briandowns/spinner"
 	"github.com/gofrs/flock"
 	homedir "github.com/mitchellh/go-homedir"
 	ini "gopkg.in/ini.v1"
@@ -73,16 +74,17 @@ type Core struct {
 
 // Config describes CLI config file and default options
 type Config struct {
-	Dir           string
-	ConfigFile    string
-	HistoryFile   string
-	LogFile       string
-	HasShell      bool
-	Core          *Core
-	ActiveProfile *ServerProfile
-	Context       *context.Context
-	Cancel        context.CancelFunc
-	C             chan bool
+	Dir            string
+	ConfigFile     string
+	HistoryFile    string
+	LogFile        string
+	HasShell       bool
+	Core           *Core
+	ActiveProfile  *ServerProfile
+	Context        *context.Context
+	Cancel         context.CancelFunc
+	C              chan bool
+	activeSpinners []*spinner.Spinner
 }
 
 // GetOutputFormats returns the supported output formats.

config/spinner.go

@@ -40,12 +40,39 @@ func (c *Config) StartSpinner(suffix string) *spinner.Spinner {
 	waiter := spinner.New(cursor, 200*time.Millisecond)
 	waiter.Suffix = " " + suffix
 	waiter.Start()
+	c.activeSpinners = append(c.activeSpinners, waiter)
 	return waiter
 }
 
 // StopSpinner stops the provided spinner if it is valid
 func (c *Config) StopSpinner(waiter *spinner.Spinner) {
 	if waiter != nil {
 		waiter.Stop()
+		for i, s := range c.activeSpinners {
+			if s == waiter {
+				c.activeSpinners = append(c.activeSpinners[:i], c.activeSpinners[i+1:]...)
+				break
+			}
+		}
+	}
+}
+
+// PauseActiveSpinners stops the spinners without removing them from the acive spinners list, allowing resume.
+func (c *Config) PauseActiveSpinners() int {
+	count := len(c.activeSpinners)
+	for _, s := range c.activeSpinners {
+		if s != nil && s.Active() {
+			s.Stop()
+		}
+	}
+	return count
+}
+
+// ResumePausedSpinners restarts the spinners from the active spinners list if they are not already running.
+func (c *Config) ResumePausedSpinners() {
+	for _, s := range c.activeSpinners {
+		if s != nil && !s.Active() {
+			s.Start()
+		}
 	}
 }