forked from projectdiscovery/shuffledns
-
Notifications
You must be signed in to change notification settings - Fork 0
/
process.go
245 lines (208 loc) · 6.4 KB
/
process.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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
package massdns
import (
"bufio"
"bytes"
"encoding/json"
"errors"
"fmt"
"os"
"os/exec"
"path/filepath"
"strconv"
"strings"
"time"
"github.com/projectdiscovery/gologger"
"github.com/mohammadanaraki/shuffledns/internal/store"
"github.com/mohammadanaraki/shuffledns/pkg/parser"
"github.com/remeh/sizedwaitgroup"
"github.com/rs/xid"
)
// Process runs the actual enumeration process returning a file
func (c *Client) Process() error {
// Process a created list or the massdns input
inputFile := c.config.InputFile
if c.config.MassdnsRaw != "" {
inputFile = c.config.MassdnsRaw
}
// Check for blank input file or non-existent input file
blank, err := IsBlankFile(inputFile)
if err != nil {
return err
}
if blank {
return errors.New("blank input file specified")
}
// Create a store for storing ip metadata
shstore := store.New()
defer shstore.Close()
// Set the correct target file
massDNSOutput := filepath.Join(c.config.TempDir, xid.New().String())
if c.config.MassdnsRaw != "" {
massDNSOutput = c.config.MassdnsRaw
}
// Check if we need to run massdns
if c.config.MassdnsRaw == "" {
// Create a temporary file for the massdns output
gologger.Info().Msgf("Creating temporary massdns output file: %s\n", massDNSOutput)
err = c.runMassDNS(massDNSOutput, shstore)
if err != nil {
return fmt.Errorf("could not execute massdns: %w", err)
}
}
gologger.Info().Msgf("Started parsing massdns output\n")
err = c.parseMassDNSOutput(massDNSOutput, shstore)
if err != nil {
return fmt.Errorf("could not parse massdns output: %w", err)
}
gologger.Info().Msgf("Massdns output parsing completed\n")
// Perform wildcard filtering only if domain name has been specified
if c.config.Domain != "" {
gologger.Info().Msgf("Started removing wildcards records\n")
err = c.filterWildcards(shstore)
if err != nil {
return fmt.Errorf("could not parse massdns output: %w", err)
}
gologger.Info().Msgf("Wildcard removal completed\n")
}
gologger.Info().Msgf("Finished enumeration, started writing output\n")
// Write the final elaborated list out
return c.writeOutput(shstore)
}
func (c *Client) runMassDNS(output string, store *store.Store) error {
if c.config.Domain != "" {
gologger.Info().Msgf("Executing massdns on %s\n", c.config.Domain)
} else {
gologger.Info().Msgf("Executing massdns\n")
}
now := time.Now()
// Run the command on a temp file and wait for the output
cmd := exec.Command(c.config.MassdnsPath, []string{"-r", c.config.ResolversFile, "-o", "Snl", "-t", "A", c.config.InputFile, "-w", output, "-s", strconv.Itoa(c.config.Threads)}...)
var stderr bytes.Buffer
cmd.Stderr = &stderr
err := cmd.Run()
if err != nil {
return fmt.Errorf("could not execute massdns: %w\ndetailed error: %s", err, stderr.String())
}
gologger.Info().Msgf("Massdns execution took %s\n", time.Since(now))
return nil
}
func (c *Client) parseMassDNSOutput(output string, store *store.Store) error {
massdnsOutput, err := os.Open(output)
if err != nil {
return fmt.Errorf("could not open massdns output file: %w", err)
}
defer massdnsOutput.Close()
// at first we need the full structure in memory to elaborate it in parallell
err = parser.Parse(massdnsOutput, func(domain string, ip []string) {
for _, ip := range ip {
// Check if ip exists in the store. If not,
// add the ip to the map and continue with the next ip.
if !store.Exists(ip) {
store.New(ip, domain)
continue
}
// Get the IP meta-information from the store.
record := store.Get(ip)
// Put the new hostname and increment the counter by 1.
record.Hostnames[domain] = struct{}{}
record.Counter++
}
})
if err != nil {
return fmt.Errorf("could not parse massdns output: %w", err)
}
return nil
}
func (c *Client) filterWildcards(st *store.Store) error {
// Start to work in parallel on wildcards
wildcardWg := sizedwaitgroup.New(c.config.WildcardsThreads)
for _, record := range st.IP {
// We've stumbled upon a wildcard, just ignore it.
c.wildcardIPMutex.Lock()
if _, ok := c.wildcardIPMap[record.IP]; ok {
c.wildcardIPMutex.Unlock()
continue
}
c.wildcardIPMutex.Unlock()
// Perform wildcard detection on the ip, if an IP is found in the wildcard
// we add it to the wildcard map so that further runs don't require such filtering again.
if record.Counter >= 5 || c.config.StrictWildcard {
wildcardWg.Add()
go func(record *store.IPMeta) {
defer wildcardWg.Done()
for host := range record.Hostnames {
isWildcard, ips := c.wildcardResolver.LookupHost(host)
if len(ips) > 0 {
c.wildcardIPMutex.Lock()
for ip := range ips {
// we add the single ip to the wildcard list
c.wildcardIPMap[ip] = struct{}{}
}
c.wildcardIPMutex.Unlock()
}
if isWildcard {
c.wildcardIPMutex.Lock()
// we also mark the original ip as wildcard, since at least once it resolved to this host
c.wildcardIPMap[record.IP] = struct{}{}
c.wildcardIPMutex.Unlock()
break
}
}
}(record)
}
}
wildcardWg.Wait()
// drop all wildcard from the store
for wildcardIP := range c.wildcardIPMap {
st.Delete(wildcardIP)
}
return nil
}
func (c *Client) writeOutput(store *store.Store) error {
// Write the unique deduplicated output to the file or stdout
// depending on what the user has asked.
var output *os.File
var w *bufio.Writer
var err error
if c.config.OutputFile != "" {
output, err = os.Create(c.config.OutputFile)
if err != nil {
return fmt.Errorf("could not create massdns output file: %v", err)
}
w = bufio.NewWriter(output)
}
buffer := &strings.Builder{}
uniqueMap := make(map[string]struct{})
for _, record := range store.IP {
for hostname := range record.Hostnames {
// Skip if we already printed this subdomain once
if _, ok := uniqueMap[hostname]; ok {
continue
}
uniqueMap[hostname] = struct{}{}
if c.config.Json {
hostnameJson, err := json.Marshal(map[string]interface{}{"hostname": hostname})
if err != nil {
return fmt.Errorf("could not marshal output as json: %v", err)
}
buffer.WriteString(string(hostnameJson))
buffer.WriteString("\n")
} else {
buffer.WriteString(hostname)
buffer.WriteString("\n")
}
data := buffer.String()
if output != nil {
_, _ = w.WriteString(data)
}
gologger.Silent().Msgf("%s", data)
buffer.Reset()
}
}
// Close the files and return
if output != nil {
w.Flush()
output.Close()
}
return nil
}