1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
|
package twofactorrecover
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"strings"
"testing"
"github.com/stretchr/testify/require"
"gitlab.com/gitlab-org/gitlab-shell/v14/client/testserver"
"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/commandargs"
"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/readwriter"
"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/twofactorrecover"
)
var (
requests []testserver.TestRequestHandler
)
func setup(t *testing.T) {
requests = []testserver.TestRequestHandler{
{
Path: "/api/v4/internal/two_factor_recovery_codes",
Handler: func(w http.ResponseWriter, r *http.Request) {
b, err := io.ReadAll(r.Body)
defer r.Body.Close()
require.NoError(t, err)
var requestBody *twofactorrecover.RequestBody
json.Unmarshal(b, &requestBody)
switch requestBody.KeyId {
case "1":
body := map[string]interface{}{
"success": true,
"recovery_codes": [2]string{"recovery", "codes"},
}
json.NewEncoder(w).Encode(body)
case "forbidden":
body := map[string]interface{}{
"success": false,
"message": "Forbidden!",
}
json.NewEncoder(w).Encode(body)
case "broken":
w.WriteHeader(http.StatusInternalServerError)
}
},
},
}
}
const (
question = "Are you sure you want to generate new two-factor recovery codes?\n" +
"Any existing recovery codes you saved will be invalidated. (yes/no)\n\n"
errorHeader = "An error occurred while trying to generate new recovery codes.\n"
)
func TestExecute(t *testing.T) {
setup(t)
url := testserver.StartSocketHttpServer(t, requests)
testCases := []struct {
desc string
arguments *commandargs.Shell
answer string
expectedOutput string
}{
{
desc: "With a known key id",
arguments: &commandargs.Shell{GitlabKeyId: "1"},
answer: "yes\n",
expectedOutput: question +
"Your two-factor authentication recovery codes are:\n\nrecovery\ncodes\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",
},
{
desc: "With bad response",
arguments: &commandargs.Shell{GitlabKeyId: "-1"},
answer: "yes\n",
expectedOutput: question + errorHeader + "Parsing failed\n",
},
{
desc: "With API returns an error",
arguments: &commandargs.Shell{GitlabKeyId: "forbidden"},
answer: "yes\n",
expectedOutput: question + errorHeader + "Forbidden!\n",
},
{
desc: "With API fails",
arguments: &commandargs.Shell{GitlabKeyId: "broken"},
answer: "yes\n",
expectedOutput: question + errorHeader + "Internal API error (500)\n",
},
{
desc: "With missing arguments",
arguments: &commandargs.Shell{},
answer: "yes\n",
expectedOutput: question + errorHeader + "who='' is invalid\n",
},
{
desc: "With negative answer",
arguments: &commandargs.Shell{},
answer: "no\n",
expectedOutput: question +
"New recovery codes have *not* been generated. Existing codes will remain valid.\n",
},
{
desc: "With some other answer",
arguments: &commandargs.Shell{},
answer: strings.Repeat("yes", 1024),
expectedOutput: question +
"New recovery codes have *not* been generated. Existing codes will remain valid.\n",
},
}
for _, tc := range testCases {
t.Run(tc.desc, func(t *testing.T) {
output := &bytes.Buffer{}
input := bytes.NewBufferString(tc.answer)
cmd := &Command{
Config: &config.Config{GitlabUrl: url},
Args: tc.arguments,
ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
}
err := cmd.Execute(context.Background())
require.NoError(t, err)
require.Equal(t, tc.expectedOutput, output.String())
})
}
}
|
