[Update] add tips for confirmation

pkg/auth/server.go

@@ -2,6 +2,7 @@ package auth
 
 import (
 	"net"
+	"strings"
 
 	"github.com/gliderlabs/ssh"
 	gossh "golang.org/x/crypto/ssh"
@@ -16,6 +17,9 @@ import (
 var mfaInstruction = "Please enter 6 digits."
 var mfaQuestion = "[MFA auth]: "
 
+var confirmInstruction = "Please wait for your admin to confirm."
+var confirmQuestion = "[YES or NO]: "
+
 const (
 	actionAccepted        = "Accepted"
 	actionFailed          = "Failed"
@@ -34,20 +38,24 @@ func checkAuth(ctx ssh.Context, password, publicKey string) (res ssh.AuthResult)
 	userClient, ok := ctx.Value(model.ContextKeyClient).(*service.SessionClient)
 	if !ok {
 		sessionClient := service.NewSessionClient(service.Username(username),
-			service.Password(password), service.PublicKey(publicKey),
 			service.RemoteAddr(remoteAddr), service.LoginType("T"))
 		userClient = &sessionClient
+		ctx.SetValue(model.ContextKeyClient, userClient)
 	}
+	userClient.SetOption(service.Password(password), service.PublicKey(publicKey))
 	user, authStatus := userClient.Authenticate(ctx)
 	switch authStatus {
 	case service.AuthMFARequired:
-		ctx.SetValue(model.ContextKeyClient, &userClient)
 		action = actionPartialAccepted
 		res = ssh.AuthPartiallySuccessful
 	case service.AuthSuccess:
 		res = ssh.AuthSuccessful
 		ctx.SetValue(model.ContextKeyUser, &user)
-
+	case service.AuthConfirmRequired:
+		required := true
+		ctx.SetValue(model.ContextKeyConfirmRequired, &required)
+		action = actionPartialAccepted
+		res = ssh.AuthPartiallySuccessful
 	default:
 		action = actionFailed
 	}
@@ -73,37 +81,64 @@ func CheckUserPublicKey(ctx ssh.Context, key ssh.PublicKey) ssh.AuthResult {
 }
 
 func CheckMFA(ctx ssh.Context, challenger gossh.KeyboardInteractiveChallenge) (res ssh.AuthResult) {
+	if value, ok := ctx.Value(model.ContextKeyConfirmFailed).(*bool); ok && *value {
+		return ssh.AuthFailed
+	}
+
 	username := ctx.User()
 	remoteAddr, _, _ := net.SplitHostPort(ctx.RemoteAddr().String())
 	res = ssh.AuthFailed
-	defer func() {
-		authMethod := "MFA"
-		if res == ssh.AuthSuccessful {
-			action := actionAccepted
-			logger.Infof("%s %s for %s from %s", action, authMethod, username, remoteAddr)
-		} else {
-			action := actionFailed
-			logger.Errorf("%s %s for %s from %s", action, authMethod, username, remoteAddr)
+
+	var confirmAction bool
+	instruction := mfaInstruction
+	question := mfaQuestion
+
+	client, ok := ctx.Value(model.ContextKeyClient).(*service.SessionClient)
+	if !ok {
+		logger.Errorf("User %s Mfa Auth failed: not found session client.", username, )
+		return
 	}
-	}()
-	answers, err := challenger(username, mfaInstruction, []string{mfaQuestion}, []bool{true})
+	value, ok := ctx.Value(model.ContextKeyConfirmRequired).(*bool)
+	if ok && *value {
+		confirmAction = true
+		instruction = confirmInstruction
+		question = confirmQuestion
+	}
+	answers, err := challenger(username, instruction, []string{question}, []bool{true})
 	if err != nil || len(answers) != 1 {
 		return
 	}
-	mfaCode := answers[0]
-	client, ok := ctx.Value(model.ContextKeyClient).(*service.SessionClient)
-	if !ok {
-		logger.Errorf("User %s Mfa Auth failed: not found session client.", username, )
+	if confirmAction {
+		switch strings.TrimSpace(strings.ToLower(answers[0])) {
+		case "yes", "y", "":
+			user, authStatus := client.CheckConfirm(ctx)
+			switch authStatus {
+			case service.AuthSuccess:
+				res = ssh.AuthSuccessful
+				ctx.SetValue(model.ContextKeyUser, &user)
 				return
 			}
+		default:
+			client.CancelConfirm()
+		}
+		failed := true
+		ctx.SetValue(model.ContextKeyConfirmFailed, &failed)
+		return
+	}
+	mfaCode := answers[0]
 	user, authStatus := client.CheckUserOTP(ctx, mfaCode)
 	switch authStatus {
 	case service.AuthSuccess:
 		res = ssh.AuthSuccessful
 		ctx.SetValue(model.ContextKeyUser, &user)
-		logger.Infof("User %s Mfa Auth success", username)
+		logger.Infof("%s MFA for %s from %s", actionAccepted, username, remoteAddr)
+	case service.AuthConfirmRequired:
+		res = ssh.AuthPartiallySuccessful
+		required := true
+		ctx.SetValue(model.ContextKeyConfirmRequired, &required)
+		logger.Infof("%s MFA for %s from %s", actionPartialAccepted, username, remoteAddr)
 	default:
-		logger.Errorf("User %s Mfa Auth failed", username)
+		logger.Errorf("%s MFA for %s from %s", actionFailed, username, remoteAddr)
 	}
 	return
 }

pkg/model/const.go

@@ -6,4 +6,6 @@ const (
 	ContextKeyUser contextKey = iota + 1
 	ContextKeyRemoteAddr
 	ContextKeyClient
+	ContextKeyConfirmRequired
+	ContextKeyConfirmFailed
 )

pkg/service/options.go

@@ -6,6 +6,7 @@ const (
 	AuthSuccess AuthStatus = iota + 1
 	AuthFailed
 	AuthMFARequired
+	AuthConfirmRequired
 )
 
 type SessionOption func(*SessionOptions)

pkg/service/users.go

@@ -55,6 +55,12 @@ type SessionClient struct {
 	authOptions map[string]AuthOptions
 }
 
+func (u *SessionClient) SetOption(setters ...SessionOption) {
+	for _, setter := range setters {
+		setter(u.option)
+	}
+}
+
 func (u *SessionClient) Authenticate(ctx context.Context) (user model.User, authStatus AuthStatus) {
 	authStatus = AuthFailed
 	data := map[string]string{
@@ -73,12 +79,8 @@ func (u *SessionClient) Authenticate(ctx context.Context) (user model.User, auth
 	if resp.Err != "" {
 		switch resp.Err {
 		case ErrLoginConfirmWait:
-			if !u.checkConfirm(ctx) {
-				logger.Errorf("User %s login confirm required err", u.option.Username)
-				return
-			}
-			logger.Infof("User %s login confirm required success", u.option.Username)
-			return u.Authenticate(ctx)
+			logger.Infof("User %s login need confirmation", u.option.Username)
+			authStatus = AuthConfirmRequired
 		case ErrMFARequired:
 			for _, item := range resp.Data.Choices {
 				u.authOptions[item] = AuthOptions{
@@ -129,16 +131,14 @@ func (u *SessionClient) CheckUserOTP(ctx context.Context, code string) (user mod
 	return
 }
 
-func (u *SessionClient) checkConfirm(ctx context.Context) (ok bool) {
+func (u *SessionClient) CheckConfirm(ctx context.Context) (user model.User, authStatus AuthStatus) {
 	var err error
+	for {
 		select {
 		case <-ctx.Done():
-		_, err = u.client.Delete(UserConfirmAuthURL, nil)
-		if err != nil {
-			logger.Errorf("User %s cancel confirmation err: %s", u.option.Username, err)
+			logger.Errorf("User %s exit and cancel confirmation", u.option.Username)
+			u.CancelConfirm()
 			return
-		}
-		logger.Infof("User %s cancel confirm request", u.option.Username)
 		case <-time.After(5 * time.Second):
 			var resp authResponse
 			_, err = u.client.Get(UserConfirmAuthURL, &resp)
@@ -150,7 +150,7 @@ func (u *SessionClient) checkConfirm(ctx context.Context) (ok bool) {
 				switch resp.Err {
 				case ErrLoginConfirmWait:
 					logger.Infof("User %s still wait confirm", u.option.Username)
-				return u.checkConfirm(ctx)
+					continue
 				case ErrLoginConfirmRejected:
 					logger.Infof("User %s confirmation was rejected by admin", u.option.Username)
 				default:
@@ -160,10 +160,19 @@ func (u *SessionClient) checkConfirm(ctx context.Context) (ok bool) {
 			}
 			if resp.Msg == "ok" {
 				logger.Infof("User %s confirmation was accepted", u.option.Username)
-			return true
+				return u.Authenticate(ctx)
+			}
 		}
 	}
+}
+
+func (u *SessionClient) CancelConfirm() {
+	_, err := u.client.Delete(UserConfirmAuthURL, nil)
+	if err != nil {
+		logger.Errorf("Cancel User %s confirmation err: %s", u.option.Username, err)
 		return
+	}
+	logger.Infof("Cancel User %s confirmation success", u.option.Username)
 }
 
 func GetUserDetail(userID string) (user *model.User) {