forked from refraction-networking/gotapdance
-
Notifications
You must be signed in to change notification settings - Fork 0
/
setup.go
228 lines (191 loc) · 6.87 KB
/
setup.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
package main
import (
"encoding/json"
"flag"
"fmt"
"net/http"
"os"
"strconv"
"strings"
"time"
atlas "github.com/keltia/ripe-atlas"
)
const (
TracerouteCost = 30
DNSCost = 10
OneOffMultiplier = 2
)
func main() {
targetsFilename := flag.String("t", "targets.json", "The JSON file from which probe targets are loaded. They can be specified by IPv4 addresses or hostnames.")
key := flag.String("key", "", "The RIPE Atlas API key used to run the measurements.")
probesFilename := flag.String("p", "probes.json", "The JSON file from which probes are loaded. See https://atlas.ripe.net/docs/api/v2/manual/measurements/types/in_detail.html for information on the format.")
arg := flag.String("arg", "", "The DNS lookup argument to use in measurements.")
verbose := flag.Bool("v", false, "Print out lots of information about RIPE Atlas API requests.")
flag.Parse()
if *key == "" {
fmt.Fprintln(os.Stderr, "You must specify an API key.")
os.Exit(1)
}
if *arg == "" {
fmt.Fprintln(os.Stderr, "You must specify a query argument.")
os.Exit(1)
}
targetsFile, err := os.Open(*targetsFilename)
if err != nil {
fmt.Fprintf(os.Stderr, `Couldn't open targets file "%s": %s\n`, *targetsFilename, err)
os.Exit(1)
}
defer targetsFile.Close()
var targets []string
err = json.NewDecoder(targetsFile).Decode(&targets)
if err != nil {
fmt.Fprintln(os.Stderr, "Couldn't decode targets file:", err)
os.Exit(1)
}
probesFile, err := os.Open(*probesFilename)
if err != nil {
fmt.Fprintf(os.Stderr, `Couldn't open probes file "%s": %s\n`, *targetsFilename, err)
os.Exit(1)
}
defer probesFile.Close()
var probes []atlas.ProbeSet
err = json.NewDecoder(probesFile).Decode(&probes)
if err != nil {
fmt.Fprintln(os.Stderr, "Couldn't decode probes file:", err)
os.Exit(1)
}
var totalProbes int
for _, p := range probes {
totalProbes += p.Requested
}
config := atlas.Config{
APIKey: *key,
Verbose: *verbose,
}
if *verbose {
config.Level = 2
}
client, err := atlas.NewClient(config)
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to create RIPE Atlas client:", err)
os.Exit(1)
}
credits, err := client.GetCredits()
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to get credits information for RIPE Atlas account:", err)
os.Exit(1)
}
creditEstimate := OneOffMultiplier * (TracerouteCost + DNSCost) * len(targets) * totalProbes
fmt.Printf("You have %d credits. By my estimation, these measurements could take up to %d credits (depending on how many probes are available). Do you want to continue? [y/N] ", credits.CurrentBalance, creditEstimate)
inputLoop:
for {
var input string
_, err = fmt.Scanf("%s", &input)
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to read input:", err)
os.Exit(1)
}
switch strings.ToLower(input) {
case "y", "yes":
break inputLoop
case "n", "no":
os.Exit(0)
default:
fmt.Printf(`I don't know what "%s" means. Please answer y or N. `, input)
}
}
fmt.Println("Creating measurements...")
dnsDefinitions := make([]atlas.Definition, 0, len(targets))
tracerouteDefinitions := make([]atlas.Definition, 0, len(targets))
tag := "refraction-routing-probe-" + time.Now().Format("2006-01-02-15-04-05")
for _, target := range targets {
dns := atlas.Definition{
Description: "DNS routing probe for " + target,
Tags: []string{tag},
Type: "dns",
AF: 4,
IsOneoff: true,
IsPublic: false,
QueryClass: "IN",
QueryType: "A",
Target: target,
QueryArgument: *arg,
}
dnsDefinitions = append(dnsDefinitions, dns)
traceroute := atlas.Definition{
Description: "Traceroute routing probe for " + target,
Tags: []string{tag},
Type: "traceroute",
AF: 4,
Protocol: "UDP",
IsOneoff: true,
IsPublic: false,
Target: target,
}
tracerouteDefinitions = append(tracerouteDefinitions, traceroute)
}
tracerouteRequest := client.NewMeasurement()
tracerouteRequest.Definitions = tracerouteDefinitions
tracerouteRequest.Probes = probes
tracerouteResp, err := client.Traceroute(tracerouteRequest)
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to create traceroute measurements:", err)
os.Exit(1)
}
fmt.Println("Successfully created traceroute measurements. Waiting for information on probes to proceed to DNS measurements.")
// The RIPE Atlas API allows us to specify another measurement ID to use the
// same probes in theory, however it seems that if the first measurement
// reports itself as failed, the second will never happen. Thus, we should manually
// get the probes to add onto the second measurement.
body := struct {
Probes []struct {
ID int `json:"id"`
} `json:"probes"`
}{}
// TODO: we can't check if body.Probes is the length we want, as some probes
// probably will not pick up the measurement. Checking for a non-zero length,
// however, isn't optimal---the probes might not be finished filling in. There
// isn't a clear way to tell whether probes are done being populated, but a
// model that waits until multiple calls end up in the same (non-zero) length
// consecutively might make more sense.
for len(body.Probes) == 0 {
// It seems like RIPE's API takes a bit of an eventual consistency model,
// so the probes aren't immediately available after creating the request.
// Keep trying until we get a non-empty probes array.
time.Sleep(10 * time.Second)
// The client doesn't allow us to get the probes of a measurement. :(
req, err := http.NewRequest("GET", fmt.Sprintf("https://atlas.ripe.net/api/v2/measurements/%d?fields=probes", tracerouteResp.Measurements[0]), nil)
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to create probes request:", err)
os.Exit(1)
}
req.Header.Add("Authorization", "Key "+*key)
resp, err := http.DefaultClient.Do(req)
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to get measurement:", err)
os.Exit(1)
}
defer resp.Body.Close()
err = json.NewDecoder(resp.Body).Decode(&body)
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to parse probes from measurement:", err)
os.Exit(1)
}
}
fmt.Println("Using probes for DNS requests:", body.Probes)
probeIDs := make([]string, 0, totalProbes)
for _, probe := range body.Probes {
probeIDs = append(probeIDs, strconv.Itoa(probe.ID))
}
probesString := strings.Join(probeIDs, ",")
dnsRequest := client.NewMeasurement()
dnsRequest.Definitions = dnsDefinitions
dnsRequest.Probes = []atlas.ProbeSet{{Requested: len(probeIDs), Type: "probes", Value: probesString}}
_, err = client.DNS(dnsRequest)
if err != nil {
fmt.Fprintln(os.Stderr, "Failed to create DNS measurements:", err)
os.Exit(1)
}
fmt.Printf("Successfully created measurements! To fetch the results, run the following command in a few minutes:\n\n")
fmt.Printf(" curl -H \"Authorization: Key %s\" https://atlas.ripe.net/api/v2/measurements/tags/%s/results/ > results.json\n", *key, tag)
}