133 lines
2.5 KiB
Go
133 lines
2.5 KiB
Go
package prober
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"dtool/utils"
|
|
)
|
|
|
|
type Data struct {
|
|
target string
|
|
dict map[string]bool
|
|
}
|
|
|
|
var dataset map[string][]string
|
|
|
|
func retrieve_ip(pool chan string, filename string) {
|
|
cnt := 0
|
|
f, err := os.Open(filename)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
fmt.Println("sending msg ...")
|
|
reader := bufio.NewReader(f)
|
|
for {
|
|
s, err := reader.ReadString('\n')
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
s = strings.Trim(s, "\n")
|
|
pool <- s
|
|
cnt++
|
|
if cnt%10 == 0 {
|
|
fmt.Println(cnt)
|
|
}
|
|
}
|
|
close(pool)
|
|
}
|
|
|
|
func active_probe(n int, addr string) Data {
|
|
target_ip := addr[:len(addr)-3]
|
|
data := Data{target_ip, make(map[string]bool)}
|
|
stop := 0
|
|
timestamp := strconv.FormatInt(time.Now().Unix(), 10)
|
|
for i := 0; i < n; i++ {
|
|
subdomain := strings.Join([]string{strings.Replace(target_ip, ".", "-", -1), "echo", strconv.Itoa(i), timestamp}, "-")
|
|
rdns_ip, err := utils.SendQuery(addr, subdomain)
|
|
if err == nil {
|
|
data.dict[rdns_ip] = true
|
|
} else {
|
|
stop += 1
|
|
}
|
|
if stop == 3 {
|
|
break
|
|
}
|
|
}
|
|
return data
|
|
}
|
|
|
|
func upstream_prober(ip_pool chan string, data_pool chan Data, wg *sync.WaitGroup) {
|
|
for {
|
|
if s, ok := <-ip_pool; ok {
|
|
addr := s + ":53"
|
|
data := active_probe(20, addr)
|
|
if data.dict != nil {
|
|
data_pool <- data
|
|
}
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
wg.Done()
|
|
}
|
|
|
|
func create_probers(num int, ip_pool chan string, data_pool chan Data, wg *sync.WaitGroup) {
|
|
for i := 0; i < num; i++ {
|
|
wg.Add(1)
|
|
go upstream_prober(ip_pool, data_pool, wg)
|
|
}
|
|
}
|
|
|
|
func store_data(pool chan Data, wg *sync.WaitGroup) {
|
|
wg.Add(1)
|
|
for {
|
|
var temp []string
|
|
if data, ok := <-pool; ok {
|
|
if len(data.dict) > 0 {
|
|
for rdns := range data.dict {
|
|
temp = append(temp, rdns)
|
|
}
|
|
dataset[data.target] = temp
|
|
}
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
wg.Done()
|
|
}
|
|
|
|
func Get_upstream_file(filename string, output string, prober_num int) {
|
|
dataset = map[string][]string{}
|
|
ip_pool := make(chan string, 500)
|
|
data_pool := make(chan Data, 200)
|
|
var probe_tasks sync.WaitGroup
|
|
var store_task sync.WaitGroup
|
|
|
|
go retrieve_ip(ip_pool, filename)
|
|
create_probers(prober_num, ip_pool, data_pool, &probe_tasks)
|
|
go store_data(data_pool, &store_task)
|
|
probe_tasks.Wait()
|
|
close(data_pool)
|
|
store_task.Wait()
|
|
utils.OutputJSON(dataset, output)
|
|
}
|
|
|
|
func Get_upstream_ip(ip string) {
|
|
dataset = make(map[string][]string)
|
|
var temp []string
|
|
addr := ip + ":53"
|
|
data := active_probe(10, addr)
|
|
for rdns := range data.dict {
|
|
temp = append(temp, rdns)
|
|
}
|
|
dataset[data.target] = temp
|
|
utils.OutputJSON(dataset, "-")
|
|
}
|