cmd.NewScanCommand - Code Metrics - Inspection of "Graceful shutdown on signint" - stefanoj3/dirstalk - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Push — master ( 6ba115...bb6332 )

by Stefano

created 2019-10-19 05:28 UTC

cmd.NewScanCommand B

↳ Parent: pkg/cmd/scan.go

Complexity

Conditions

Size

Total Lines	98
Code Lines	70

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
eloc	70
nop	1
dl	0
loc	98
rs	7.9818
c	0
b	0
f	0

How to fix Long Method

package cmd

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"os/signal"

	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/stefanoj3/dirstalk/pkg/cmd/termination"
	"github.com/stefanoj3/dirstalk/pkg/common"
	"github.com/stefanoj3/dirstalk/pkg/dictionary"
	"github.com/stefanoj3/dirstalk/pkg/scan"
	"github.com/stefanoj3/dirstalk/pkg/scan/client"
	"github.com/stefanoj3/dirstalk/pkg/scan/filter"
	"github.com/stefanoj3/dirstalk/pkg/scan/output"
	"github.com/stefanoj3/dirstalk/pkg/scan/producer"
	"github.com/stefanoj3/dirstalk/pkg/scan/summarizer"
	"github.com/stefanoj3/dirstalk/pkg/scan/summarizer/tree"
)

func NewScanCommand(logger *logrus.Logger) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "scan [url]",
		Short: "Scan the given URL",
		RunE:  buildScanFunction(logger),
	}

	cmd.Flags().StringP(
		flagScanDictionary,
		flagScanDictionaryShort,
		"",
		"dictionary to use for the scan (path to local file or remote url)",
	)
	common.Must(cmd.MarkFlagFilename(flagScanDictionary))
	common.Must(cmd.MarkFlagRequired(flagScanDictionary))

	cmd.Flags().StringSlice(
		flagScanHTTPMethods,
		[]string{"GET"},
		"comma separated list of http methods to use; eg: GET,POST,PUT",
	)

	cmd.Flags().IntSlice(
		flagScanHTTPStatusesToIgnore,
		[]int{http.StatusNotFound},
		"comma separated list of http statuses to ignore when showing and processing results; eg: 404,301",
	)

	cmd.Flags().IntP(
		flagScanThreads,
		flagScanThreadsShort,
		3,
		"amount of threads for concurrent requests",
	)

	cmd.Flags().IntP(
		flagScanHTTPTimeout,
		"",
		5000,
		"timeout in milliseconds",
	)

	cmd.Flags().BoolP(
		flagScanHTTPCacheRequests,
		"",
		true,
		"cache requests to avoid performing the same request multiple times within the same scan (EG if the "+
			"server reply with the same redirect location multiple times, dirstalk will follow it only once)",
	)

	cmd.Flags().IntP(
		flagScanScanDepth,
		"",
		3,
		"scan depth",
	)

	cmd.Flags().StringP(
		flagScanSocks5Host,
		"",
		"",
		"socks5 host to use",
	)

	cmd.Flags().StringP(
		flagScanUserAgent,
		"",
		"",
		"user agent to use for http requests",
	)

	cmd.Flags().BoolP(
		flagScanCookieJar,
		"",
		false,
		"enables the use of a cookie jar: it will retain any cookie sent "+
			"from the server and send them for the following requests",
	)

	cmd.Flags().StringArray(
		flagScanCookie,
		[]string{},
		"cookie to add to each request; eg name=value (can be specified multiple times)",
	)

	cmd.Flags().StringArray(
		flagScanHeader,
		[]string{},
		"header to add to each request; eg name=value (can be specified multiple times)",
	)

	cmd.Flags().String(
		flagScanResultOutput,
		"",
		"path where to store result output",
	)

	return cmd
}

func buildScanFunction(logger *logrus.Logger) func(cmd *cobra.Command, args []string) error {
	f := func(cmd *cobra.Command, args []string) error {
		u, err := getURL(args)
		if err != nil {
			return err
		}

		cnf, err := scanConfigFromCmd(cmd)
		if err != nil {
			return errors.Wrap(err, "failed to build config")
		}

		return startScan(logger, cnf, u)
	}

	return f
}

func getURL(args []string) (*url.URL, error) {
	if len(args) == 0 {
		return nil, errors.New("no URL provided")
	}

	arg := args[0]

	u, err := url.ParseRequestURI(arg)
	if err != nil {
		return nil, errors.Wrap(err, "the first argument must be a valid url")
	}

	return u, nil
}

// startScan is a convenience method that wires together all the dependencies needed to start a scan
func startScan(logger *logrus.Logger, cnf *scan.Config, u *url.URL) error {
	c, err := client.NewClientFromConfig(
		cnf.TimeoutInMilliseconds,
		cnf.Socks5Url,
		cnf.UserAgent,
		cnf.UseCookieJar,
		cnf.Cookies,
		cnf.Headers,
		cnf.CacheRequests,
		u,
	)
	if err != nil {
		return errors.Wrap(err, "failed to build client")
	}

	dict, err := dictionary.NewDictionaryFrom(cnf.DictionaryPath, c)
	if err != nil {
		return errors.Wrap(err, "failed to build dictionary")
	}

	targetProducer := producer.NewDictionaryProducer(cnf.HTTPMethods, dict, cnf.ScanDepth)
	reproducer := producer.NewReProducer(targetProducer)

	resultFilter := filter.NewHTTPStatusResultFilter(cnf.HTTPStatusesToIgnore)

	s := scan.NewScanner(
		c,
		targetProducer,
		reproducer,
		resultFilter,
		logger,
	)

	logger.WithFields(logrus.Fields{
		"url":               u.String(),
		"threads":           cnf.Threads,
		"dictionary-length": len(dict),
		"scan-depth":        cnf.ScanDepth,
		"timeout":           cnf.TimeoutInMilliseconds,
		"socks5":            cnf.Socks5Url,
		"cookies":           stringifyCookies(cnf.Cookies),
		"cookie-jar":        cnf.UseCookieJar,
		"headers":           stringifyHeaders(cnf.Headers),
		"user-agent":        cnf.UserAgent,
	}).Info("Starting scan")

	resultSummarizer := summarizer.NewResultSummarizer(tree.NewResultTreeProducer(), logger)

	osSigint := make(chan os.Signal, 1)
	signal.Notify(osSigint, os.Interrupt)

	outputSaver, err := newOutputSaver(cnf.Out)
	if err != nil {
		return errors.Wrap(err, "failed to create output saver")
	}

	defer func() {
		resultSummarizer.Summarize()
		err := outputSaver.Close()
		if err != nil {
			logger.WithError(err).Error("failed to close output file")
		}
		logger.Info("Finished scan")
	}()

	ctx, cancellationFunc := context.WithCancel(context.Background())
	defer cancellationFunc()

	resultsChannel := s.Scan(ctx, u, cnf.Threads)

	terminationHandler := termination.NewTerminationHandler(2)

	for {
		select {
		case <-osSigint:
			terminationHandler.SignalTermination()
			cancellationFunc()

			if terminationHandler.ShouldTerminate() {
				logger.Info("Received sigint, terminating...")
				return nil
			}

			logger.Info(
				"Received sigint, trying to shutdown gracefully, another SIGNINT will terminate the application",
			)
		case result, ok := <-resultsChannel:
			if !ok {
				logger.Debug("result channel is being closed, scan should be complete")
				return nil
			}

			resultSummarizer.Add(result)

			if err := outputSaver.Save(result); err != nil {
				return errors.Wrap(err, "failed to add output to file")
			}
		}
	}
}

func newOutputSaver(path string) (OutputSaver, error) {
	if path == "" {
		return output.NewNullSaver(), nil
	}

	return output.NewFileSaver(path)
}

func stringifyCookies(cookies []*http.Cookie) string {
	result := ""

	for _, cookie := range cookies {
		result += fmt.Sprintf("{%s=%s}", cookie.Name, cookie.Value)
	}

	return result
}

func stringifyHeaders(headers map[string]string) string {
	result := ""

	for name, value := range headers {
		result += fmt.Sprintf("{%s:%s}", name, value)
	}

	return result
}


1			package cmd
2
3			import (
4			"context"
5			"fmt"
6			"net/http"
7			"net/url"
8			"os"
9			"os/signal"
10
11			"github.com/pkg/errors"
12			"github.com/sirupsen/logrus"
13			"github.com/spf13/cobra"
14			"github.com/stefanoj3/dirstalk/pkg/cmd/termination"
15			"github.com/stefanoj3/dirstalk/pkg/common"
16			"github.com/stefanoj3/dirstalk/pkg/dictionary"
17			"github.com/stefanoj3/dirstalk/pkg/scan"
18			"github.com/stefanoj3/dirstalk/pkg/scan/client"
19			"github.com/stefanoj3/dirstalk/pkg/scan/filter"
20			"github.com/stefanoj3/dirstalk/pkg/scan/output"
21			"github.com/stefanoj3/dirstalk/pkg/scan/producer"
22			"github.com/stefanoj3/dirstalk/pkg/scan/summarizer"
23			"github.com/stefanoj3/dirstalk/pkg/scan/summarizer/tree"
24			)
25
26			func NewScanCommand(logger logrus.Logger) cobra.Command {
27			cmd := &cobra.Command{
28			Use: "scan [url]",
29			Short: "Scan the given URL",
30			RunE: buildScanFunction(logger),
31			}
32
33			cmd.Flags().StringP(
34			flagScanDictionary,
35			flagScanDictionaryShort,
36			"",
37			"dictionary to use for the scan (path to local file or remote url)",
38			)
39			common.Must(cmd.MarkFlagFilename(flagScanDictionary))
40			common.Must(cmd.MarkFlagRequired(flagScanDictionary))
41
42			cmd.Flags().StringSlice(
43			flagScanHTTPMethods,
44			[]string{"GET"},
45			"comma separated list of http methods to use; eg: GET,POST,PUT",
46			)
47
48			cmd.Flags().IntSlice(
49			flagScanHTTPStatusesToIgnore,
50			[]int{http.StatusNotFound},
51			"comma separated list of http statuses to ignore when showing and processing results; eg: 404,301",
52			)
53
54			cmd.Flags().IntP(
55			flagScanThreads,
56			flagScanThreadsShort,
57			3,
58			"amount of threads for concurrent requests",
59			)
60
61			cmd.Flags().IntP(
62			flagScanHTTPTimeout,
63			"",
64			5000,
65			"timeout in milliseconds",
66			)
67
68			cmd.Flags().BoolP(
69			flagScanHTTPCacheRequests,
70			"",
71			true,
72			"cache requests to avoid performing the same request multiple times within the same scan (EG if the "+
73			"server reply with the same redirect location multiple times, dirstalk will follow it only once)",
74			)
75
76			cmd.Flags().IntP(
77			flagScanScanDepth,
78			"",
79			3,
80			"scan depth",
81			)
82
83			cmd.Flags().StringP(
84			flagScanSocks5Host,
85			"",
86			"",
87			"socks5 host to use",
88			)
89
90			cmd.Flags().StringP(
91			flagScanUserAgent,
92			"",
93			"",
94			"user agent to use for http requests",
95			)
96
97			cmd.Flags().BoolP(
98			flagScanCookieJar,
99			"",
100			false,
101			"enables the use of a cookie jar: it will retain any cookie sent "+
102			"from the server and send them for the following requests",
103			)
104
105			cmd.Flags().StringArray(
106			flagScanCookie,
107			[]string{},
108			"cookie to add to each request; eg name=value (can be specified multiple times)",
109			)
110
111			cmd.Flags().StringArray(
112			flagScanHeader,
113			[]string{},
114			"header to add to each request; eg name=value (can be specified multiple times)",
115			)
116
117			cmd.Flags().String(
118			flagScanResultOutput,
119			"",
120			"path where to store result output",
121			)
122
123			return cmd
124			}
125
126			func buildScanFunction(logger logrus.Logger) func(cmd cobra.Command, args []string) error {
127			f := func(cmd *cobra.Command, args []string) error {
128			u, err := getURL(args)
129			if err != nil {
130			return err
131			}
132
133			cnf, err := scanConfigFromCmd(cmd)
134			if err != nil {
135			return errors.Wrap(err, "failed to build config")
136			}
137
138			return startScan(logger, cnf, u)
139			}
140
141			return f
142			}
143
144			func getURL(args []string) (*url.URL, error) {
145			if len(args) == 0 {
146			return nil, errors.New("no URL provided")
147			}
148
149			arg := args[0]
150
151			u, err := url.ParseRequestURI(arg)
152			if err != nil {
153			return nil, errors.Wrap(err, "the first argument must be a valid url")
154			}
155
156			return u, nil
157			}
158
159			// startScan is a convenience method that wires together all the dependencies needed to start a scan
160			func startScan(logger logrus.Logger, cnf scan.Config, u *url.URL) error {
161			c, err := client.NewClientFromConfig(
162			cnf.TimeoutInMilliseconds,
163			cnf.Socks5Url,
164			cnf.UserAgent,
165			cnf.UseCookieJar,
166			cnf.Cookies,
167			cnf.Headers,
168			cnf.CacheRequests,
169			u,
170			)
171			if err != nil {
172			return errors.Wrap(err, "failed to build client")
173			}
174
175			dict, err := dictionary.NewDictionaryFrom(cnf.DictionaryPath, c)
176			if err != nil {
177			return errors.Wrap(err, "failed to build dictionary")
178			}
179
180			targetProducer := producer.NewDictionaryProducer(cnf.HTTPMethods, dict, cnf.ScanDepth)
181			reproducer := producer.NewReProducer(targetProducer)
182
183			resultFilter := filter.NewHTTPStatusResultFilter(cnf.HTTPStatusesToIgnore)
184
185			s := scan.NewScanner(
186			c,
187			targetProducer,
188			reproducer,
189			resultFilter,
190			logger,
191			)
192
193			logger.WithFields(logrus.Fields{
194			"url": u.String(),
195			"threads": cnf.Threads,
196			"dictionary-length": len(dict),
197			"scan-depth": cnf.ScanDepth,
198			"timeout": cnf.TimeoutInMilliseconds,
199			"socks5": cnf.Socks5Url,
200			"cookies": stringifyCookies(cnf.Cookies),
201			"cookie-jar": cnf.UseCookieJar,
202			"headers": stringifyHeaders(cnf.Headers),
203			"user-agent": cnf.UserAgent,
204			}).Info("Starting scan")
205
206			resultSummarizer := summarizer.NewResultSummarizer(tree.NewResultTreeProducer(), logger)
207
208			osSigint := make(chan os.Signal, 1)
209			signal.Notify(osSigint, os.Interrupt)
210
211			outputSaver, err := newOutputSaver(cnf.Out)
212			if err != nil {
213			return errors.Wrap(err, "failed to create output saver")
214			}
215
216			defer func() {
217			resultSummarizer.Summarize()
218			err := outputSaver.Close()
219			if err != nil {
220			logger.WithError(err).Error("failed to close output file")
221			}
222			logger.Info("Finished scan")
223			}()
224
225			ctx, cancellationFunc := context.WithCancel(context.Background())
226			defer cancellationFunc()
227
228			resultsChannel := s.Scan(ctx, u, cnf.Threads)
229
230			terminationHandler := termination.NewTerminationHandler(2)
231
232			for {
233			select {
234			case <-osSigint:
235			terminationHandler.SignalTermination()
236			cancellationFunc()
237
238			if terminationHandler.ShouldTerminate() {
239			logger.Info("Received sigint, terminating...")
240			return nil
241			}
242
243			logger.Info(
244			"Received sigint, trying to shutdown gracefully, another SIGNINT will terminate the application",
245			)
246			case result, ok := <-resultsChannel:
247			if !ok {
248			logger.Debug("result channel is being closed, scan should be complete")
249			return nil
250			}
251
252			resultSummarizer.Add(result)
253
254			if err := outputSaver.Save(result); err != nil {
255			return errors.Wrap(err, "failed to add output to file")
256			}
257			}
258			}
259			}
260
261			func newOutputSaver(path string) (OutputSaver, error) {
262			if path == "" {
263			return output.NewNullSaver(), nil
264			}
265
266			return output.NewFileSaver(path)
267			}
268
269			func stringifyCookies(cookies []*http.Cookie) string {
270			result := ""
271
272			for _, cookie := range cookies {
273			result += fmt.Sprintf("{%s=%s}", cookie.Name, cookie.Value)
274			}
275
276			return result
277			}
278
279			func stringifyHeaders(headers map[string]string) string {
280			result := ""
281
282			for name, value := range headers {
283			result += fmt.Sprintf("{%s:%s}", name, value)
284			}
285
286			return result
287			}
288

stefanoj3 / dirstalk

Push — master ( 6ba115...bb6332 )

cmd.NewScanCommand B

Complexity

Size

Duplication

Importance

How to fix Long Method

Long Method

Duplication Side-by-Side

Filter issues like