diff options
Diffstat (limited to 'internal/gitlabnet')
| -rw-r--r-- | internal/gitlabnet/twofactorverify/client.go | 90 | ||||
| -rw-r--r-- | internal/gitlabnet/twofactorverify/client_test.go | 154 |
2 files changed, 244 insertions, 0 deletions
diff --git a/internal/gitlabnet/twofactorverify/client.go b/internal/gitlabnet/twofactorverify/client.go new file mode 100644 index 0000000..aab302b --- /dev/null +++ b/internal/gitlabnet/twofactorverify/client.go @@ -0,0 +1,90 @@ +package twofactorverify + +import ( + "context" + "errors" + "fmt" + "net/http" + + "gitlab.com/gitlab-org/gitlab-shell/client" + "gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs" + "gitlab.com/gitlab-org/gitlab-shell/internal/config" + "gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet" + "gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/discover" +) + +type Client struct { + config *config.Config + client *client.GitlabNetClient +} + +type Response struct { + Success bool `json:"success"` + Message string `json:"message"` +} + +type RequestBody struct { + KeyId string `json:"key_id,omitempty"` + UserId int64 `json:"user_id,omitempty"` + OTPAttempt string `json:"otp_attempt"` +} + +func NewClient(config *config.Config) (*Client, error) { + client, err := gitlabnet.GetClient(config) + if err != nil { + return nil, fmt.Errorf("Error creating http client: %v", err) + } + + 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) + if err != nil { + return err + } + + response, err := c.client.Post(ctx, "/two_factor_otp_check", requestBody) + if err != nil { + return err + } + defer response.Body.Close() + + return parse(response) +} + +func parse(hr *http.Response) error { + response := &Response{} + if err := gitlabnet.ParseJSON(hr, response); err != nil { + return err + } + + if !response.Success { + return errors.New(response.Message) + } + + return nil +} + +func (c *Client) getRequestBody(ctx context.Context, args *commandargs.Shell, otp string) (*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, OTPAttempt: otp} + } else { + userInfo, err := client.GetByCommandArgs(ctx, args) + + if err != nil { + return nil, err + } + + requestBody = &RequestBody{UserId: userInfo.UserId, OTPAttempt: otp} + } + + return requestBody, nil +} diff --git a/internal/gitlabnet/twofactorverify/client_test.go b/internal/gitlabnet/twofactorverify/client_test.go new file mode 100644 index 0000000..7bb037e --- /dev/null +++ b/internal/gitlabnet/twofactorverify/client_test.go @@ -0,0 +1,154 @@ +package twofactorverify + +import ( + "context" + "encoding/json" + "gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/discover" + "io/ioutil" + "net/http" + "testing" + + "github.com/stretchr/testify/require" + "gitlab.com/gitlab-org/gitlab-shell/client" + "gitlab.com/gitlab-org/gitlab-shell/client/testserver" + "gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs" + "gitlab.com/gitlab-org/gitlab-shell/internal/config" +) + +func initialize(t *testing.T) []testserver.TestRequestHandler { + requests := []testserver.TestRequestHandler{ + { + Path: "/api/v4/internal/two_factor_otp_check", + Handler: func(w http.ResponseWriter, r *http.Request) { + b, err := ioutil.ReadAll(r.Body) + defer r.Body.Close() + + require.NoError(t, err) + + var requestBody *RequestBody + require.NoError(t, json.Unmarshal(b, &requestBody)) + + switch requestBody.KeyId { + case "0": + body := map[string]interface{}{ + "success": true, + } + require.NoError(t, json.NewEncoder(w).Encode(body)) + case "1": + body := map[string]interface{}{ + "success": false, + "message": "error message", + } + require.NoError(t, json.NewEncoder(w).Encode(body)) + case "2": + w.WriteHeader(http.StatusForbidden) + body := &client.ErrorResponse{ + Message: "Not allowed!", + } + require.NoError(t, json.NewEncoder(w).Encode(body)) + case "3": + w.Write([]byte("{ \"message\": \"broken json!\"")) + case "4": + w.WriteHeader(http.StatusForbidden) + } + + if requestBody.UserId == 1 { + body := map[string]interface{}{ + "success": true, + } + require.NoError(t, 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", + } + require.NoError(t, json.NewEncoder(w).Encode(body)) + }, + }, + } + + return requests +} + +const ( + otpAttempt = "123456" +) + +func TestVerifyOTPByKeyId(t *testing.T) { + client, cleanup := setup(t) + defer cleanup() + + args := &commandargs.Shell{GitlabKeyId: "0"} + err := client.VerifyOTP(context.Background(), args, otpAttempt) + require.NoError(t, err) +} + +func TestVerifyOTPByUsername(t *testing.T) { + client, cleanup := setup(t) + defer cleanup() + + args := &commandargs.Shell{GitlabUsername: "jane-doe"} + err := client.VerifyOTP(context.Background(), args, otpAttempt) + require.NoError(t, err) +} + +func TestErrorMessage(t *testing.T) { + client, cleanup := setup(t) + defer cleanup() + + args := &commandargs.Shell{GitlabKeyId: "1"} + err := client.VerifyOTP(context.Background(), args, otpAttempt) + require.Equal(t, "error message", err.Error()) +} + +func TestErrorResponses(t *testing.T) { + client, cleanup := setup(t) + defer cleanup() + + testCases := []struct { + desc string + fakeId string + expectedError string + }{ + { + desc: "A response with an error message", + fakeId: "2", + expectedError: "Not allowed!", + }, + { + desc: "A response with bad JSON", + fakeId: "3", + expectedError: "Parsing failed", + }, + { + desc: "An error response without message", + fakeId: "4", + expectedError: "Internal API error (403)", + }, + } + + for _, tc := range testCases { + t.Run(tc.desc, func(t *testing.T) { + args := &commandargs.Shell{GitlabKeyId: tc.fakeId} + err := client.VerifyOTP(context.Background(), args, otpAttempt) + + require.EqualError(t, err, tc.expectedError) + }) + } +} + +func setup(t *testing.T) (*Client, func()) { + requests := initialize(t) + url, cleanup := testserver.StartSocketHttpServer(t, requests) + + client, err := NewClient(&config.Config{GitlabUrl: url}) + require.NoError(t, err) + + return client, cleanup +} |