forked from sensepost/gowitness
-
Notifications
You must be signed in to change notification settings - Fork 0
/
file.go
146 lines (120 loc) · 3.64 KB
/
file.go
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
144
145
146
package cmd
import (
"bufio"
"net/url"
"os"
"strings"
"github.com/remeh/sizedwaitgroup"
"github.com/sensepost/gowitness/lib"
"github.com/spf13/cobra"
)
// fileCmd represents the file command
var fileCmd = &cobra.Command{
Use: "file [input]",
Short: "screenshot URLs sourced from a file or stdin",
Long: `Screenshot URLs sourced from a file or stdin. URLs in the source
file should be newline separated. Invalid URLs are simply logged and ignored.`,
Example: `$ gowitness file -f ~/Desktop/urls
$ gowitness file -f urls.txt --threads 2
$ cat urls.txt | gowitness file -f -
$ gowitness file -f <( shuf domains ) --no-http`,
Run: func(cmd *cobra.Command, args []string) {
log := options.Logger
scanner, f, err := getScanner(options.File)
if err != nil {
log.Fatal().Err(err).Str("file", options.File).Msg("unable to read source file")
}
defer f.Close()
db, err := db.Get()
if err != nil {
log.Fatal().Err(err).Msg("failed to get a db handle")
}
log.Debug().Int("threads", options.Threads).Msg("thread count to use with goroutines")
swg := sizedwaitgroup.New(options.Threads)
if err = options.PrepareScreenshotPath(); err != nil {
log.Fatal().Err(err).Msg("failed to prepare the screenshot path")
}
for scanner.Scan() {
candidate := scanner.Text()
if candidate == "" {
return
}
for _, u := range getUrls(candidate) {
swg.Add()
log.Debug().Str("url", u.String()).Msg("queueing goroutine for url")
go func(url *url.URL) {
defer swg.Done()
p := &lib.Processor{
Logger: log,
Db: db,
Chrome: chrm,
URL: url,
ScreenshotPath: options.ScreenshotPath,
}
if err := p.Gowitness(); err != nil {
log.Error().Err(err).Str("url", url.String()).Msg("failed to witness url")
}
}(u)
}
}
swg.Wait()
log.Info().Msg("processing complete")
},
}
func init() {
rootCmd.AddCommand(fileCmd)
fileCmd.Flags().StringVarP(&options.File, "file", "f", "", "file containing urls. use - for stdin")
fileCmd.Flags().IntVarP(&options.Threads, "threads", "t", 4, "threads used to run")
fileCmd.Flags().BoolVar(&options.NoHTTPS, "no-https", false, "do not prefix https:// where missing")
fileCmd.Flags().BoolVar(&options.NoHTTP, "no-http", false, "do not prefix http:// where missing")
cobra.MarkFlagRequired(fileCmd.Flags(), "file")
}
// getInput determines what the file input should be
// without any file argument we will assume stdin with -
func getInput(a []string) (input string) {
if len(a) <= 0 {
input = "-"
} else {
input = a[0]
}
return
}
// getScanner prepares a bufio.Scanner to read from either
// stdin, or a file.
// the size attribute > 0 will be returned if a file was the input
// it is up to the caller to close the file.
func getScanner(i string) (*bufio.Scanner, *os.File, error) {
if i == "-" {
return bufio.NewScanner(os.Stdin), nil, nil
}
file, err := os.Open(i)
if err != nil {
return nil, nil, err
}
return bufio.NewScanner(file), file, nil
}
// getUrls generates urls for an incoming target depending
// on wether the target has an http prefix and the flags set
func getUrls(target string) (c []*url.URL) {
// if there already is a protocol, just parse and add that
if strings.HasPrefix(target, "http") {
u, err := url.Parse(target)
if err == nil {
c = append(c, u)
}
return
}
if !strings.HasPrefix(target, "http://") && !options.NoHTTP {
u, err := url.Parse("http://" + target)
if err == nil {
c = append(c, u)
}
}
if !strings.HasPrefix(target, "https://") && !options.NoHTTPS {
u, err := url.Parse("https://" + target)
if err == nil {
c = append(c, u)
}
}
return
}