forked from d3mondev/puredns
-
Notifications
You must be signed in to change notification settings - Fork 0
/
callback.go
160 lines (129 loc) · 3.35 KB
/
callback.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
package massdns
import (
"fmt"
"os"
"strings"
)
type state int
const (
stateNewAnswerSection state = iota
stateSaveAnswer
stateSkip
)
// DefaultWriteCallback is a callback that can save massdns results to files.
// It can save the valid domains found and the massdns results that gave valid domains.
type DefaultWriteCallback struct {
massdnsFile *os.File
domainFile *os.File
curState state
curDomain string
domainSaved bool
found int
}
var _ Callback = (*DefaultWriteCallback)(nil)
// NewDefaultWriteCallback creates a new DefaultWriteCallback.
// The file names can be empty to disable saving to a file.
func NewDefaultWriteCallback(massdnsFilename string, domainFilename string) (*DefaultWriteCallback, error) {
cb := &DefaultWriteCallback{}
// Create a writer that writes massdns answers
if massdnsFilename != "" {
file, err := os.Create(massdnsFilename)
if err != nil {
return nil, err
}
cb.massdnsFile = file
}
// Create a writer that writes valid domains found
if domainFilename != "" {
file, err := os.Create(domainFilename)
if err != nil {
return nil, err
}
cb.domainFile = file
}
return cb, nil
}
// Callback reads a line from the massdns stdout handler, parses the output and
// saves the relevant data.
func (c *DefaultWriteCallback) Callback(line string) error {
// Don't parse JSON if we're not saving anything
if c.domainFile == nil && c.massdnsFile == nil {
return nil
}
// If we receive an empty line, it's the start of a new answer
if line == "" {
c.curState = stateNewAnswerSection
return nil
}
switch c.curState {
// We're at the beginning of a new answer section, look for the domain name
case stateNewAnswerSection:
parts := strings.Split(line, " ")
if len(parts) != 3 {
c.curState = stateSkip
return nil
}
domain := strings.TrimSuffix(parts[0], ".")
if domain == "" {
c.curState = stateSkip
return nil
}
c.curDomain = domain
c.curState = stateSaveAnswer
c.domainSaved = false
fallthrough
// Save the answer record found
case stateSaveAnswer:
parts := strings.Split(line, " ")
if len(parts) != 3 {
c.curState = stateSkip
return nil
}
domain := c.curDomain
rrType := parts[1]
answer := strings.TrimSuffix(parts[2], ".")
// Only look for A, AAAA, and CNAME records
if rrType != "A" && rrType != "AAAA" && rrType != "CNAME" {
return nil
}
// If we haven't saved the domain yet, save it
if !c.domainSaved {
c.saveDomain(c.curDomain)
c.domainSaved = true
c.found++
}
// Valid record found, save it
return c.saveLine(fmt.Sprintf("%s %s %s", domain, rrType, answer))
// Answer was invalid, skip until we receive a new answer section
case stateSkip:
return nil
}
return nil
}
// saveLine saves a line to the massdns file.
func (c *DefaultWriteCallback) saveLine(line string) error {
if c.massdnsFile != nil {
_, err := c.massdnsFile.WriteString(line + "\n")
return err
}
return nil
}
// saveDomain saves a domain to the domain file.
func (c *DefaultWriteCallback) saveDomain(domain string) error {
if c.domainFile != nil {
_, err := c.domainFile.WriteString(domain + "\n")
return err
}
return nil
}
// Close closes the writers.
func (c *DefaultWriteCallback) Close() {
if c.massdnsFile != nil {
c.massdnsFile.Sync()
c.massdnsFile.Close()
}
if c.domainFile != nil {
c.domainFile.Sync()
c.domainFile.Close()
}
}