Add prototype code written by Fortinet

2 files changed, 99 insertions, 23 deletions

diff --git a/internal/command/twofactorverify/twofactorverify.go b/internal/command/twofactorverify/twofactorverify.go
index b1c5508..02b5ac8 100644
--- a/internal/command/twofactorverify/twofactorverify.go
+++ b/internal/command/twofactorverify/twofactorverify.go
@@ -17,12 +17,44 @@ type Command struct {
 	ReadWriter *readwriter.ReadWriter
 }
 
+type Result struct {
+	Error   error
+	Status  string
+	Success bool
+}
+
 func (c *Command) Execute(ctx context.Context) error {
-	err := c.verifyOTP(ctx, c.getOTP())
-	if err != nil {
-		return err
-	}
+	verify := make(chan Result)
+	pushauth := make(chan Result)
+
+	go func() {
+		status, success, err := c.verifyOTP(ctx, c.getOTP())
+		verify <- Result{Error: err, Status: status, Success: success}
+	}()
 
+	go func() {
+		status, success, err := c.pushAuth(ctx)
+		pushauth <- Result{Error: err, Status: status, Success: success}
+	}()
+
+L:
+	for {
+		select {
+		case res := <-verify:
+			if res.Error != nil {
+				return res.Error
+			}
+			fmt.Fprint(c.ReadWriter.Out, res.Status)
+			break L
+		case res := <-pushauth:
+			if res.Success {
+				fmt.Fprint(c.ReadWriter.Out, res.Status)
+				break L
+			} else {
+				// ignore reject from remote, need to wait for user input in this case
+			}
+		}
+	}
 	return nil
 }
 
@@ -38,18 +70,46 @@ func (c *Command) getOTP() string {
 	return answer
 }
 
-func (c *Command) verifyOTP(ctx context.Context, otp string) error {
+func (c *Command) pushAuth(ctx context.Context) (status string, success bool, err error) {
 	client, err := twofactorverify.NewClient(c.Config)
 	if err != nil {
-		return err
+		return "", false, err
 	}
 
-	err = client.VerifyOTP(ctx, c.Args, otp)
-	if err == nil {
-		fmt.Fprint(c.ReadWriter.Out, "\nOTP validation successful. Git operations are now allowed.\n")
+	reason := ""
+
+	success, reason, err = client.PushAuth(ctx, c.Args)
+	if success {
+		status = fmt.Sprintf("\nPush OTP validation successful. Git operations are now allowed.\n")
 	} else {
-		fmt.Fprintf(c.ReadWriter.Out, "\nOTP validation failed.\n%v\n", err)
+		if err != nil {
+			status = fmt.Sprintf("\nPush OTP validation failed.\n%v\n", err)
+		} else {
+			status = fmt.Sprintf("\nPush OTP validation failed.\n%v\n", reason)
+		}
 	}
 
-	return nil
+	return
+}
+
+func (c *Command) verifyOTP(ctx context.Context, otp string) (status string, success bool, err error) {
+	client, err := twofactorverify.NewClient(c.Config)
+	if err != nil {
+		return "", false, err
+	}
+
+	reason := ""
+
+	success, reason, err = client.VerifyOTP(ctx, c.Args, otp)
+	if success {
+		status = fmt.Sprintf("\nOTP validation successful. Git operations are now allowed.\n")
+	} else {
+		if err != nil {
+			status = fmt.Sprintf("\nOTP validation failed.\n%v\n", err)
+		} else {
+			status = fmt.Sprintf("\nOTP validation failed.\n%v\n", reason)
+		}
+	}
+
+	return
 }
diff --git a/internal/gitlabnet/twofactorverify/client.go b/internal/gitlabnet/twofactorverify/client.go
index aab302b..cf22cf8 100644
--- a/internal/gitlabnet/twofactorverify/client.go
+++ b/internal/gitlabnet/twofactorverify/client.go
@@ -2,7 +2,6 @@ package twofactorverify
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"net/http"
 
@@ -27,6 +26,7 @@ type RequestBody struct {
 	KeyId      string `json:"key_id,omitempty"`
 	UserId     int64  `json:"user_id,omitempty"`
 	OTPAttempt string `json:"otp_attempt"`
+	PushAuth   bool   `json:"push_auth"`
 }
 
 func NewClient(config *config.Config) (*Client, error) {
@@ -38,35 +38,51 @@ func NewClient(config *config.Config) (*Client, error) {
 	return &Client{config: config, client: client}, nil
 }
 
-func (c *Client) VerifyOTP(ctx context.Context, args *commandargs.Shell, otp string) error {
-	requestBody, err := c.getRequestBody(ctx, args, otp)
+func (c *Client) VerifyOTP(ctx context.Context, args *commandargs.Shell, otp string) (bool, string, error) {
+	requestBody, err := c.getRequestBody(ctx, args, otp, false)
 	if err != nil {
-		return err
+		return false, "", err
 	}
 
 	response, err := c.client.Post(ctx, "/two_factor_otp_check", requestBody)
 	if err != nil {
-		return err
+		return false, "", err
 	}
 	defer response.Body.Close()
 
 	return parse(response)
 }
 
-func parse(hr *http.Response) error {
+func (c *Client) PushAuth(ctx context.Context, args *commandargs.Shell) (bool, string, error) {
+	// enable push auth in internal rest api
+	requestBody, err := c.getRequestBody(ctx, args, "", true)
+	if err != nil {
+		return false, "", err
+	}
+
+	response, err := c.client.Post(ctx, "/two_factor_otp_check", requestBody)
+	if err != nil {
+		return false, "", err
+	}
+	defer response.Body.Close()
+
+	return parse(response)
+}
+
+func parse(hr *http.Response) (bool, string, error) {
 	response := &Response{}
 	if err := gitlabnet.ParseJSON(hr, response); err != nil {
-		return err
+		return false, "", err
 	}
 
 	if !response.Success {
-		return errors.New(response.Message)
+		return false, response.Message, nil
 	}
 
-	return nil
+	return true, response.Message, nil
 }
 
-func (c *Client) getRequestBody(ctx context.Context, args *commandargs.Shell, otp string) (*RequestBody, error) {
+func (c *Client) getRequestBody(ctx context.Context, args *commandargs.Shell, otp string, pushauth bool) (*RequestBody, error) {
 	client, err := discover.NewClient(c.config)
 
 	if err != nil {
@@ -75,7 +91,7 @@ func (c *Client) getRequestBody(ctx context.Context, args *commandargs.Shell, ot
 
 	var requestBody *RequestBody
 	if args.GitlabKeyId != "" {
-		requestBody = &RequestBody{KeyId: args.GitlabKeyId, OTPAttempt: otp}
+		requestBody = &RequestBody{KeyId: args.GitlabKeyId, OTPAttempt: otp, PushAuth: pushauth}
 	} else {
 		userInfo, err := client.GetByCommandArgs(ctx, args)
 
@@ -83,7 +99,7 @@ func (c *Client) getRequestBody(ctx context.Context, args *commandargs.Shell, ot
 			return nil, err
 		}
 
-		requestBody = &RequestBody{UserId: userInfo.UserId, OTPAttempt: otp}
+		requestBody = &RequestBody{UserId: userInfo.UserId, OTPAttempt: otp, PushAuth: pushauth}
 	}
 
 	return requestBody, nil