From 17b3c808ca5c2ffc7522905c23b1bba2a092321c Mon Sep 17 00:00:00 2001 From: Igor Drozdov Date: Wed, 20 Mar 2019 23:35:11 +0300 Subject: Implement getting recovery codes by username --- go/internal/command/discover/discover.go | 10 +-- .../command/twofactorrecover/twofactorrecover.go | 6 +- .../twofactorrecover/twofactorrecover_test.go | 2 +- go/internal/gitlabnet/discover/client.go | 13 ++++ go/internal/gitlabnet/twofactorrecover/client.go | 39 +++++++++-- .../gitlabnet/twofactorrecover/client_test.go | 42 ++++++++++-- spec/gitlab_shell_two_factor_recovery_spec.rb | 77 ++++++++++++++-------- 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 -- cgit v1.2.1