-
Notifications
You must be signed in to change notification settings - Fork 2.3k
/
dsl.go
140 lines (124 loc) · 3.48 KB
/
dsl.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
package dsl
import (
"fmt"
"strings"
"github.com/Knetic/govaluate"
"github.com/miekg/dns"
"github.com/projectdiscovery/dsl"
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/dns/dnsclientpool"
"github.com/projectdiscovery/nuclei/v3/pkg/types"
sliceutil "github.com/projectdiscovery/utils/slice"
stringsutil "github.com/projectdiscovery/utils/strings"
)
var (
HelperFunctions map[string]govaluate.ExpressionFunction
FunctionNames []string
// knownPorts is a list of known ports for protocols implemented in nuclei
knowPorts = []string{"80", "443", "8080", "8081", "8443", "53"}
)
func init() {
_ = dsl.AddFunction(dsl.NewWithMultipleSignatures("resolve", []string{
"(host string) string",
"(format string) string",
}, false, func(args ...interface{}) (interface{}, error) {
argCount := len(args)
if argCount == 0 || argCount > 2 {
return nil, dsl.ErrInvalidDslFunction
}
format := "4"
var dnsType uint16
if len(args) > 1 {
format = strings.ToLower(types.ToString(args[1]))
}
switch format {
case "4", "a":
dnsType = dns.TypeA
case "6", "aaaa":
dnsType = dns.TypeAAAA
case "cname":
dnsType = dns.TypeCNAME
case "ns":
dnsType = dns.TypeNS
case "txt":
dnsType = dns.TypeTXT
case "srv":
dnsType = dns.TypeSRV
case "ptr":
dnsType = dns.TypePTR
case "mx":
dnsType = dns.TypeMX
case "soa":
dnsType = dns.TypeSOA
case "caa":
dnsType = dns.TypeCAA
default:
return nil, fmt.Errorf("invalid dns type")
}
err := dnsclientpool.Init(&types.Options{})
if err != nil {
return nil, err
}
dnsClient, err := dnsclientpool.Get(nil, &dnsclientpool.Configuration{})
if err != nil {
return nil, err
}
// query
rawResp, err := dnsClient.Query(types.ToString(args[0]), dnsType)
if err != nil {
return nil, err
}
dnsValues := map[uint16][]string{
dns.TypeA: rawResp.A,
dns.TypeAAAA: rawResp.AAAA,
dns.TypeCNAME: rawResp.CNAME,
dns.TypeNS: rawResp.NS,
dns.TypeTXT: rawResp.TXT,
dns.TypeSRV: rawResp.SRV,
dns.TypePTR: rawResp.PTR,
dns.TypeMX: rawResp.MX,
dns.TypeCAA: rawResp.CAA,
dns.TypeSOA: rawResp.GetSOARecords(),
}
if values, ok := dnsValues[dnsType]; ok {
firstFound, found := sliceutil.FirstNonZero(values)
if found {
return firstFound, nil
}
}
return "", fmt.Errorf("no records found")
}))
_ = dsl.AddFunction(dsl.NewWithMultipleSignatures("getNetworkPort", []string{
"(Port string,defaultPort string) string)",
"(Port int,defaultPort int) int",
}, false, func(args ...interface{}) (interface{}, error) {
if len(args) != 2 {
return nil, dsl.ErrInvalidDslFunction
}
port := types.ToString(args[0])
defaultPort := types.ToString(args[1])
if port == "" || stringsutil.EqualFoldAny(port, knowPorts...) {
return defaultPort, nil
}
return port, nil
}))
dsl.PrintDebugCallback = func(args ...interface{}) error {
gologger.Info().Msgf("print_debug value: %s", fmt.Sprint(args))
return nil
}
HelperFunctions = dsl.HelperFunctions()
FunctionNames = dsl.GetFunctionNames(HelperFunctions)
}
type CompilationError struct {
DslSignature string
WrappedError error
}
func (e *CompilationError) Error() string {
return fmt.Sprintf("could not compile DSL expression %q: %v", e.DslSignature, e.WrappedError)
}
func (e *CompilationError) Unwrap() error {
return e.WrappedError
}
func GetPrintableDslFunctionSignatures(noColor bool) string {
return dsl.GetPrintableDslFunctionSignatures(noColor)
}