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

7 files changed, 140 insertions, 49 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)
diff --git a/spec/gitlab_shell_two_factor_recovery_spec.rb b/spec/gitlab_shell_two_factor_recovery_spec.rb
index 872fa85..19999e5 100644
--- a/spec/gitlab_shell_two_factor_recovery_spec.rb
+++ b/spec/gitlab_shell_two_factor_recovery_spec.rb
@@ -10,7 +10,12 @@ describe 'bin/gitlab-shell 2fa_recovery_codes' do
       res.content_type = 'application/json'
       res.status = 200
 
-      key_id = req.query['key_id'] || JSON.parse(req.body)['key_id']
+      key_id = req.query['key_id'] || req.query['user_id']
+
+      unless key_id
+        body = JSON.parse(req.body)
+        key_id = body['key_id'] || body['user_id'].to_s
+      end
 
       if key_id == '100'
         res.body = '{"success":true, "recovery_codes": ["1", "2"]}'
@@ -18,43 +23,63 @@ describe 'bin/gitlab-shell 2fa_recovery_codes' do
         res.body = '{"success":false, "message": "Forbidden!"}'
       end
     end
+
+    server.mount_proc('/api/v4/internal/discover') do |req, res|
+      res.status = 200
+      res.content_type = 'application/json'
+      res.body = '{"id":100, "name": "Some User", "username": "someuser"}'
+    end
   end
 
   shared_examples 'dialog for regenerating recovery keys' do
-    context 'when runs successfully' do
-      let(:cmd) { "#{gitlab_shell_path} key-100" }
+    context 'when the user agrees to regenerate keys' do
+      def verify_successful_regeneration!(cmd)
+        Open3.popen2(env, cmd) do |stdin, stdout|
+          expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
+          expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+
+          stdin.puts('yes')
+
+          expect(stdout.flush.read).to eq(
+            "\nYour two-factor authentication recovery codes are:\n\n" \
+            "1\n2\n\n" \
+            "During sign in, use one of the codes above when prompted for\n" \
+            "your two-factor code. Then, visit your Profile Settings and add\n" \
+            "a new device so you do not lose access to your account again.\n"
+          )
+        end
+      end
+
+      context 'when key is provided' do
+        let(:cmd) { "#{gitlab_shell_path} key-100" }
 
-      context 'when the user agrees to regenerate keys' do
         it 'the recovery keys are regenerated' do
-          Open3.popen2(env, cmd) do |stdin, stdout|
-            expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
-            expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+          verify_successful_regeneration!(cmd)
+        end
+      end
 
-            stdin.puts('yes')
+      context 'when username is provided' do
+        let(:cmd) { "#{gitlab_shell_path} username-someone" }
 
-            expect(stdout.flush.read).to eq(
-              "\nYour two-factor authentication recovery codes are:\n\n" \
-              "1\n2\n\n" \
-              "During sign in, use one of the codes above when prompted for\n" \
-              "your two-factor code. Then, visit your Profile Settings and add\n" \
-              "a new device so you do not lose access to your account again.\n"
-            )
-          end
+        it 'the recovery keys are regenerated' do
+          verify_successful_regeneration!(cmd)
         end
       end
+    end
 
-      context 'when the user disagrees to regenerate keys' do
-        it 'the recovery keys are not regenerated' do
-          Open3.popen2(env, cmd) do |stdin, stdout|
-            expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
-            expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+    context 'when the user disagrees to regenerate keys' do
+      let(:cmd) { "#{gitlab_shell_path} key-100" }
 
-            stdin.puts('no')
+      it 'the recovery keys are not regenerated' do
+        Open3.popen2(env, cmd) do |stdin, stdout|
+          expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
+          expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
 
-            expect(stdout.flush.read).to eq(
-              "\nNew recovery codes have *not* been generated. Existing codes will remain valid.\n"
-            )
-          end
+          stdin.puts('no')
+
+          expect(stdout.flush.read).to eq(
+            "\nNew recovery codes have *not* been generated. Existing codes will remain valid.\n"
+          )
         end
       end
     end