Implement getting recovery codes by usernameid-go-recovery-codes

6 files changed, 89 insertions, 23 deletions

diff --git a/go/internal/command/discover/discover.go b/go/internal/command/discover/discover.go
index 9b53a51..9bb442f 100644
--- a/go/internal/command/discover/discover.go
+++ b/go/internal/command/discover/discover.go
@@ -35,13 +35,5 @@ func (c *Command) getUserInfo() (*discover.Response, error) {
 		return nil, err
 	}
 
-	if c.Args.GitlabKeyId != "" {
-		return client.GetByKeyId(c.Args.GitlabKeyId)
-	} else if c.Args.GitlabUsername != "" {
-		return client.GetByUsername(c.Args.GitlabUsername)
-	} else {
-		// There was no 'who' information, this  matches the ruby error
-		// message.
-		return nil, fmt.Errorf("who='' is invalid")
-	}
+	return client.GetByCommandArgs(c.Args)
 }
diff --git a/go/internal/command/twofactorrecover/twofactorrecover.go b/go/internal/command/twofactorrecover/twofactorrecover.go
index 227f462..e77a334 100644
--- a/go/internal/command/twofactorrecover/twofactorrecover.go
+++ b/go/internal/command/twofactorrecover/twofactorrecover.go
@@ -54,15 +54,11 @@ func (c *Command) displayRecoveryCodes(readWriter *readwriter.ReadWriter) {
 }
 
 func (c *Command) getRecoveryCodes() ([]string, error) {
-	if c.Args.GitlabKeyId == "" {
-		return nil, fmt.Errorf("Failed to get key id")
-	}
-
 	client, err := twofactorrecover.NewClient(c.Config)
 
 	if err != nil {
 		return nil, err
 	}
 
-	return client.GetRecoveryCodes(c.Args.GitlabKeyId)
+	return client.GetRecoveryCodes(c.Args)
 }
diff --git a/go/internal/command/twofactorrecover/twofactorrecover_test.go b/go/internal/command/twofactorrecover/twofactorrecover_test.go
index a005af9..908ee13 100644
--- a/go/internal/command/twofactorrecover/twofactorrecover_test.go
+++ b/go/internal/command/twofactorrecover/twofactorrecover_test.go
@@ -108,7 +108,7 @@ func TestExecute(t *testing.T) {
 			desc:           "With missing arguments",
 			arguments:      &commandargs.CommandArgs{},
 			answer:         "yes\n",
-			expectedOutput: question + errorHeader + "Failed to get key id\n",
+			expectedOutput: question + errorHeader + "who='' is invalid\n",
 		},
 		{
 			desc:      "With negative answer",
diff --git a/go/internal/gitlabnet/discover/client.go b/go/internal/gitlabnet/discover/client.go
index 8df78fb..e84b1b4 100644
--- a/go/internal/gitlabnet/discover/client.go
+++ b/go/internal/gitlabnet/discover/client.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet"
 )
@@ -30,6 +31,18 @@ func NewClient(config *config.Config) (*Client, error) {
 	return &Client{config: config, client: client}, nil
 }
 
+func (c *Client) GetByCommandArgs(args *commandargs.CommandArgs) (*Response, error) {
+	if args.GitlabKeyId != "" {
+		return c.GetByKeyId(args.GitlabKeyId)
+	} else if args.GitlabUsername != "" {
+		return c.GetByUsername(args.GitlabUsername)
+	} else {
+		// There was no 'who' information, this  matches the ruby error
+		// message.
+		return nil, fmt.Errorf("who='' is invalid")
+	}
+}
+
 func (c *Client) GetByKeyId(keyId string) (*Response, error) {
 	params := url.Values{}
 	params.Add("key_id", keyId)
diff --git a/go/internal/gitlabnet/twofactorrecover/client.go b/go/internal/gitlabnet/twofactorrecover/client.go
index 606e742..2e47c64 100644
--- a/go/internal/gitlabnet/twofactorrecover/client.go
+++ b/go/internal/gitlabnet/twofactorrecover/client.go
@@ -7,8 +7,10 @@ import (
 	"io/ioutil"
 	"net/http"
 
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/discover"
 )
 
 type Client struct {
@@ -23,7 +25,8 @@ type Response struct {
 }
 
 type RequestBody struct {
-	KeyId string `json:"key_id"`
+	KeyId  string `json:"key_id,omitempty"`
+	UserId int64  `json:"user_id,omitempty"`
 }
 
 func NewClient(config *config.Config) (*Client, error) {
@@ -35,9 +38,14 @@ func NewClient(config *config.Config) (*Client, error) {
 	return &Client{config: config, client: client}, nil
 }
 
-func (c *Client) GetRecoveryCodes(gitlabKeyId string) ([]string, error) {
-	values := RequestBody{KeyId: gitlabKeyId}
-	response, err := c.client.Post("/two_factor_recovery_codes", values)
+func (c *Client) GetRecoveryCodes(args *commandargs.CommandArgs) ([]string, error) {
+	requestBody, err := c.getRequestBody(args)
+
+	if err != nil {
+		return nil, err
+	}
+
+	response, err := c.client.Post("/two_factor_recovery_codes", requestBody)
 
 	if err != nil {
 		return nil, err
@@ -71,3 +79,26 @@ func (c *Client) parseResponse(resp *http.Response) (*Response, error) {
 		return parsedResponse, nil
 	}
 }
+
+func (c *Client) getRequestBody(args *commandargs.CommandArgs) (*RequestBody, error) {
+	client, err := discover.NewClient(c.config)
+
+	if err != nil {
+		return nil, err
+	}
+
+	var requestBody *RequestBody
+	if args.GitlabKeyId != "" {
+		requestBody = &RequestBody{KeyId: args.GitlabKeyId}
+	} else {
+		userInfo, err := client.GetByCommandArgs(args)
+
+		if err != nil {
+			return nil, err
+		}
+
+		requestBody = &RequestBody{UserId: userInfo.UserId}
+	}
+
+	return requestBody, nil
+}
diff --git a/go/internal/gitlabnet/twofactorrecover/client_test.go b/go/internal/gitlabnet/twofactorrecover/client_test.go
index 7b0b189..5cbc011 100644
--- a/go/internal/gitlabnet/twofactorrecover/client_test.go
+++ b/go/internal/gitlabnet/twofactorrecover/client_test.go
@@ -9,8 +9,10 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/discover"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/testserver"
 )
 
@@ -57,25 +59,56 @@ func initialize(t *testing.T) {
 				case "4":
 					w.WriteHeader(http.StatusForbidden)
 				}
+
+				if requestBody.UserId == 1 {
+					body := map[string]interface{}{
+						"success":        true,
+						"recovery_codes": [2]string{"recovery 2", "codes 2"},
+					}
+					json.NewEncoder(w).Encode(body)
+				}
+			},
+		},
+		{
+			Path: "/api/v4/internal/discover",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				body := &discover.Response{
+					UserId:   1,
+					Username: "jane-doe",
+					Name:     "Jane Doe",
+				}
+				json.NewEncoder(w).Encode(body)
 			},
 		},
 	}
 }
 
-func TestGetRecoveryCodes(t *testing.T) {
+func TestGetRecoveryCodesByKeyId(t *testing.T) {
 	client, cleanup := setup(t)
 	defer cleanup()
 
-	result, err := client.GetRecoveryCodes("0")
+	args := &commandargs.CommandArgs{GitlabKeyId: "0"}
+	result, err := client.GetRecoveryCodes(args)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"recovery 1", "codes 1"}, result)
 }
 
+func TestGetRecoveryCodesByUsername(t *testing.T) {
+	client, cleanup := setup(t)
+	defer cleanup()
+
+	args := &commandargs.CommandArgs{GitlabUsername: "jane-doe"}
+	result, err := client.GetRecoveryCodes(args)
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"recovery 2", "codes 2"}, result)
+}
+
 func TestMissingUser(t *testing.T) {
 	client, cleanup := setup(t)
 	defer cleanup()
 
-	_, err := client.GetRecoveryCodes("1")
+	args := &commandargs.CommandArgs{GitlabKeyId: "1"}
+	_, err := client.GetRecoveryCodes(args)
 	assert.Equal(t, "missing user", err.Error())
 }
 
@@ -107,7 +140,8 @@ func TestErrorResponses(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			resp, err := client.GetRecoveryCodes(tc.fakeId)
+			args := &commandargs.CommandArgs{GitlabKeyId: tc.fakeId}
+			resp, err := client.GetRecoveryCodes(args)
 
 			assert.EqualError(t, err, tc.expectedError)
 			assert.Nil(t, resp)