Passed
Pull Request — master (#30)
by Stefano
02:32
created

scan_test.TestShouldUseTheSpecifiedUserAgent   A

Complexity

Conditions 4

Size

Total Lines 37
Code Lines 27

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 4
eloc 27
nop 1
dl 0
loc 37
rs 9.232
c 0
b 0
f 0
1
package scan_test
2
3
import (
4
	"fmt"
5
	"net"
6
	"net/http"
7
	"net/http/httptest"
8
	"net/url"
9
	"sync"
10
	"testing"
11
	"time"
12
13
	socks5 "github.com/armon/go-socks5"
14
	"github.com/chuckpreslar/emission"
15
	"github.com/stefanoj3/dirstalk/pkg/common/test"
16
	"github.com/stefanoj3/dirstalk/pkg/scan"
17
	"github.com/stretchr/testify/assert"
18
)
19
20
const socks5TestServerHost = "127.0.0.1:8899"
21
22
func TestStartScan(t *testing.T) {
23
	logger, _ := test.NewLogger()
24
25
	requestMap := &sync.Map{}
26
	testServer := buildTestServer(requestMap)
27
	defer testServer.Close()
28
29
	u, err := url.Parse(testServer.URL)
30
	assert.NoError(t, err)
31
32
	config := &scan.Config{
33
		Threads:               3,
34
		DictionaryPath:        "testdata/dictionary1.txt",
35
		HTTPMethods:           []string{http.MethodGet},
36
		TimeoutInMilliseconds: 50,
37
		ScanDepth:             3,
38
	}
39
40
	actualResults := make([]scan.Target, 0, 2)
41
	mx := sync.Mutex{}
42
	eventManager := emission.NewEmitter()
43
	eventManager.On(scan.EventResultFound, func(r *scan.Result) {
44
		if r.Response.StatusCode != 200 {
45
			return
46
		}
47
48
		mx.Lock()
49
		defer mx.Unlock()
50
		actualResults = append(actualResults, r.Target)
51
	})
52
53
	err = scan.StartScan(logger, eventManager, config, u)
54
	assert.NoError(t, err)
55
56
	// Asserting which requests are made to the remote service
57
	assertRequest(t, true, http.MethodGet, "/home", requestMap)
58
	assertRequest(t, true, http.MethodGet, "/about", requestMap)
59
	assertRequest(t, true, http.MethodGet, "/home/home", requestMap)
60
	assertRequest(t, true, http.MethodGet, "/home/about", requestMap)
61
	assertRequest(t, true, http.MethodGet, "/home/about/home", requestMap)
62
	assertRequest(t, true, http.MethodGet, "/home/about/about", requestMap)
63
64
	assertRequest(t, false, http.MethodDelete, "/home", requestMap)
65
	assertRequest(t, false, http.MethodDelete, "/about", requestMap)
66
67
	assertRequest(t, false, http.MethodPatch, "/home", requestMap)
68
	assertRequest(t, false, http.MethodPatch, "/about", requestMap)
69
70
	assertRequest(t, false, http.MethodPost, "/home", requestMap)
71
	assertRequest(t, false, http.MethodPost, "/about", requestMap)
72
73
	assertRequest(t, false, http.MethodPut, "/home", requestMap)
74
	assertRequest(t, false, http.MethodPut, "/about", requestMap)
75
76
	assertRequest(t, false, http.MethodGet, "/about/home", requestMap)
77
	assertRequest(t, false, http.MethodGet, "/about/about", requestMap)
78
	// -----------------------------------------------------------
79
80
	// Asserting on the actual results found - considering only 200s for this test
81
	assert.Len(t, actualResults, 2)
82
	assertTargetsContains(t, scan.Target{Depth: 3, Path: "/home", Method: http.MethodGet}, actualResults)
83
	assertTargetsContains(t, scan.Target{Depth: 2, Path: "/home/about", Method: http.MethodGet}, actualResults)
84
}
85
86
func TestStartScanWithSocks5ShouldFindResultsWhenAServerIsAvailable(t *testing.T) {
87
	logger, _ := test.NewLogger()
88
89
	testServer := buildTestServer(&sync.Map{})
90
	defer testServer.Close()
91
92
	u, err := url.Parse(testServer.URL)
93
	assert.NoError(t, err)
94
95
	socks5URL, err := url.Parse("socks5://" + socks5TestServerHost)
96
	assert.NoError(t, err)
97
98
	config := &scan.Config{
99
		Threads:               3,
100
		DictionaryPath:        "testdata/dictionary1.txt",
101
		HTTPMethods:           []string{http.MethodGet},
102
		TimeoutInMilliseconds: 50,
103
		ScanDepth:             3,
104
		Socks5Url:             socks5URL,
105
	}
106
107
	listener := startSocks5TestServer(t)
108
	defer listener.Close()
109
110
	actualResults := make([]scan.Target, 0, 2)
111
	mx := sync.Mutex{}
112
	eventManager := emission.NewEmitter()
113
	eventManager.On(scan.EventResultFound, func(r *scan.Result) {
114
		if r.Response.StatusCode != 200 {
115
			return
116
		}
117
118
		mx.Lock()
119
		defer mx.Unlock()
120
		actualResults = append(actualResults, r.Target)
121
	})
122
123
	err = scan.StartScan(logger, eventManager, config, u)
124
	assert.NoError(t, err)
125
126
	assert.Len(t, actualResults, 2)
127
	assertTargetsContains(t, scan.Target{Depth: 3, Path: "/home", Method: http.MethodGet}, actualResults)
128
	assertTargetsContains(t, scan.Target{Depth: 2, Path: "/home/about", Method: http.MethodGet}, actualResults)
129
}
130
131
func TestShouldUseTheSpecifiedUserAgent(t *testing.T) {
132
	const testUserAgent = "my_test_user_agent"
133
134
	logger, _ := test.NewLogger()
135
136
	var request *http.Request
137
	doneChannel := make(chan bool)
138
139
	testServer := httptest.NewServer(
140
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
141
			request = r
142
			doneChannel <- true
143
		}),
144
	)
145
	defer testServer.Close()
146
147
	u, err := url.Parse(testServer.URL)
148
	assert.NoError(t, err)
149
150
	config := &scan.Config{
151
		Threads:               3,
152
		DictionaryPath:        "testdata/one_element_dictionary.txt",
153
		HTTPMethods:           []string{http.MethodGet},
154
		TimeoutInMilliseconds: 50,
155
		ScanDepth:             3,
156
		UserAgent:             testUserAgent,
157
	}
158
159
	eventManager := emission.NewEmitter()
160
	err = scan.StartScan(logger, eventManager, config, u)
161
	assert.NoError(t, err)
162
163
	select {
164
	case <-doneChannel:
165
		assert.Equal(t, testUserAgent, request.Header.Get("User-Agent"))
166
	case <-time.After(time.Second * 1):
167
		t.Fatal("failed to receive request")
168
	}
169
}
170
171
func TestShouldFailToScanWithAnUnreachableSocks5Server(t *testing.T) {
172
	logger, loggerBuffer := test.NewLogger()
173
174
	requestMap := &sync.Map{}
175
	testServer := buildTestServer(requestMap)
176
	defer testServer.Close()
177
178
	u, err := url.Parse(testServer.URL)
179
	assert.NoError(t, err)
180
181
	socks5URL, err := url.Parse("socks5://127.0.0.1:12345")
182
	assert.NoError(t, err)
183
184
	config := &scan.Config{
185
		Threads:               3,
186
		DictionaryPath:        "testdata/dictionary2.txt",
187
		HTTPMethods:           []string{http.MethodGet},
188
		TimeoutInMilliseconds: 50,
189
		ScanDepth:             3,
190
		Socks5Url:             socks5URL,
191
	}
192
193
	listener := startSocks5TestServer(t)
194
	defer listener.Close()
195
196
	actualResults := make([]scan.Target, 0, 2)
197
	mx := sync.Mutex{}
198
	eventManager := emission.NewEmitter()
199
	eventManager.On(scan.EventResultFound, func(r *scan.Result) {
200
		if r.Response.StatusCode != 200 {
201
			return
202
		}
203
204
		mx.Lock()
205
		defer mx.Unlock()
206
		actualResults = append(actualResults, r.Target)
207
	})
208
209
	err = scan.StartScan(logger, eventManager, config, u)
210
	assert.NoError(t, err)
211
212
	assert.Len(t, actualResults, 0)
213
	assert.Contains(t, loggerBuffer.String(), "connection refused")
214
215
	requestMap.Range(func(key, value interface{}) bool {
216
		t.Fatal("no request was supposed to be recorded: socks5 is down, the server should remain unreachable")
217
		return true
218
	})
219
}
220
221
func assertTargetsContains(t *testing.T, target scan.Target, results []scan.Target) {
222
	for _, actualResult := range results {
223
		if target == actualResult {
224
			return
225
		}
226
	}
227
228
	t.Errorf("Target %v not found in %v", target, results)
229
}
230
231
func assertRequest(t *testing.T, expected bool, method, path string, requestMap *sync.Map) {
232
	if expected {
233
		assert.True(
234
			t,
235
			hasRequest(requestMap, method, path),
236
			"expected request for `%s %s`, none received",
237
			method,
238
			path,
239
		)
240
	} else {
241
		assert.False(
242
			t,
243
			hasRequest(requestMap, method, path),
244
			"no request was expected for `%s %s`",
245
			method,
246
			path,
247
		)
248
	}
249
}
250
251
func buildTestServer(requestMap *sync.Map) *httptest.Server {
252
	return httptest.NewServer(
253
		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
254
			storeRequest(requestMap, r.Method, r.URL.Path)
255
256
			if r.Method == http.MethodGet && r.URL.Path == "/home" {
257
				w.WriteHeader(200)
258
				return
259
			}
260
261
			if r.Method == http.MethodGet && r.URL.Path == "/home/about" {
262
				w.WriteHeader(200)
263
				return
264
			}
265
266
			w.WriteHeader(http.StatusNotFound)
267
		}),
268
	)
269
}
270
271
func startSocks5TestServer(t *testing.T) net.Listener {
272
	conf := &socks5.Config{}
273
	server, err := socks5.New(conf)
274
	if err != nil {
275
		t.Fatalf("failed to create socks5: %s", err.Error())
276
	}
277
278
	listener, err := net.Listen("tcp", socks5TestServerHost)
279
	if err != nil {
280
		t.Fatalf("failed to create listener: %s", err.Error())
281
	}
282
283
	go func() {
284
		// Create SOCKS5 proxy on localhost port 8000
285
		if err := server.Serve(listener); err != nil {
286
			t.Logf("socks5 stopped serving: %s", err.Error())
287
		}
288
	}()
289
290
	return listener
291
}
292
293
func storeRequest(requestMap *sync.Map, method, path string) {
294
	requestMap.Store(methodAndPathToString(method, path), true)
295
}
296
297
func hasRequest(requestMap *sync.Map, method, path string) bool {
298
	_, ok := requestMap.Load(methodAndPathToString(method, path))
299
	return ok
300
}
301
302
func methodAndPathToString(method, path string) string {
303
	return fmt.Sprintf("%s_%s", method, path)
304
}
305