initial commit

prober/rdns_prober.go

@@ -0,0 +1,133 @@
+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 {
+		fmt.Printf("cannot open file %s\n", filename)
+		return
+	}
+	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, 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)
+}
+
+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)
+}