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
|
package cmd
import (
"dtool/utils"
"fmt"
"strconv"
"strings"
"sync"
"time"
"github.com/spf13/cobra"
)
// input parameters
// var query_target string
var query_port string
var query_name string
var query_type uint16
var query_class uint16
var query_rate int
var query_packets int
var query_random_prefix bool
/*
DNS请求测试功能，可设置参数包括域名（随机前缀），资源记录type/class，请求速率，请求数量
*/
var queryCmd = &cobra.Command{
Use: "query",
Short: "",
Long: "",
Run: queryTest,
}
func concurrent_request(ip, domain, port string, qtype, qclass uint16, result chan int, wg *sync.WaitGroup) {
wg.Add(1)
res, err := utils.SendQuery(ip, domain, port, qtype, qclass)
if err != nil {
result <- (-1)
//fmt.Println(err)
} else {
rcode := utils.ParseRcode(res)
result <- rcode
}
wg.Done()
}
func response_stat(stat []int, result chan int, wg *sync.WaitGroup) {
wg.Add(1)
defer wg.Done()
for {
if v, ok := <-result; ok {
switch v {
case 0:
stat[0]++
case 2:
stat[1]++
case 3:
stat[2]++
case -1:
stat[3]++
default:
stat[4]++
}
} else {
break
}
}
}
func queryTest(cmd *cobra.Command, args []string) {
var wg1 sync.WaitGroup
var wg2 sync.WaitGroup
result := make(chan int)
stat := make([]int, 5)
if len(args) == 1 {
if utils.IsValidIP(args[0]) {
timestamp := strconv.FormatInt(time.Now().Unix(), 10)
go response_stat(stat, result, &wg2)
for i := 0; i < query_packets; i++ {
var qname string
if query_random_prefix {
qname = strings.Join([]string{timestamp, strconv.Itoa(i)}, "-") + query_name
} else {
qname = query_name
}
go concurrent_request(args[0], qname, query_port, query_type, query_class, result, &wg1)
time.Sleep(time.Second / time.Duration(query_rate))
}
wg1.Wait()
close(result)
wg2.Wait()
fmt.Printf("total: %v noerr: %v servfail: %v nx: %v error: %v other: %v\n", query_packets, stat[0], stat[1], stat[2], stat[3], stat[4])
}
}
}
func init() {
queryCmd.Flags().BoolVarP(&query_random_prefix, "random", "R", false, "random query name")
queryCmd.Flags().StringVarP(&query_port, "port", "p", "53", "query target port")
queryCmd.Flags().StringVarP(&query_name, "name", "d", ".example.com", "query name in DNS requests")
queryCmd.Flags().Uint16VarP(&query_type, "type", "t", 1, "target record type")
queryCmd.Flags().Uint16VarP(&query_class, "class", "c", 1, "target record class")
queryCmd.Flags().IntVarP(&query_rate, "rate", "r", 10, "request rate (packets per second)")
queryCmd.Flags().IntVarP(&query_packets, "packets", "n", 100, "total request number")
rootCmd.AddCommand(queryCmd)
}
|
