package common

import (
	"fmt"
	"time"

	"github.com/miekg/dns"
)

type Pair struct {
	Nameserver string
	Domain     string
}

var (
	Nameservers, Valid, Recursive []string
	Found                         bool
)

func message(domain string, reqtype uint16, ra bool) *dns.Msg {
	msg := new(dns.Msg)
	msg.Id = dns.Id()
	msg.RecursionDesired = ra
	msg.Question = make([]dns.Question, 1)
	msg.Question[0] = dns.Question{dns.Fqdn(domain), reqtype, dns.ClassINET}
	return msg
}

func ParseNS(nservers []string) ([]string, []string) {
	var valid, recursive []string
	msg := message("supernets.org", dns.TypeA, false)
	for _, ns := range nservers {
		in, err := dns.Exchange(msg, ns+":53")
		if err != nil {
			Error("nameserver " + ns + " is not responding")
			continue
		}
		if in.Rcode == dns.RcodeRefused {
			Warning("nameserver " + ns + " refused the test query, non-recursive snooping may not be viable")
		}
		if in.RecursionAvailable {
			Success("nameserver " + ns + " is recursive")
			recursive = append(recursive, ns)
		}
		valid = append(valid, ns)
	}
	return valid, recursive
}

func TestReq() bool {
	msg := message("cloudflare.com", dns.TypeA, false)
	in, err := dns.Exchange(msg, "1.1.1.1:53")
	if err != nil {
		return false
	}
	if len(in.Answer) > 0 {
		return true
	}
	return false
}

func PullNS(d string) {
	nsmsg := message(d, dns.TypeNS, true)
	in, err := dns.Exchange(nsmsg, "1.1.1.1:53")
	if err != nil {
		Fatal("unable to retrieve nameservers for " + d)
	}

	for _, ans := range in.Answer {
		ns, ok := ans.(*dns.NS)
		if ok {
			Nameservers = append(Nameservers, ns.Ns)
		}
	}

}

func Verify() {
	if !TestReq() {
		Error("neutral non-recursive query was refused, are you on a vpn or dirty box?")
	}
	Success("neutral non-recursive test query succeeded")

	Valid, Recursive = ParseNS(Nameservers)
	Info(fmt.Sprintf("%d/%d nameservers are recursive", len(Recursive), len(Valid)))

	if len(Valid) == 0 {
		Fatal("no valid nameservers available")
	}
}

func Query(q <-chan Pair, tracker chan<- interface{}, delay int) {
	for pair := range q {
		msg := message(pair.Domain, dns.TypeA, false)
		in, err := dns.Exchange(msg, pair.Nameserver+":53")
		if err != nil {
			Error(err.Error())
			continue
		}

		if len(in.Answer) > 0 {
			Found = true
			fmt.Printf("[%s] associated domain %s found on %s\n", Vendors[Domains[pair.Domain].Vendor], pair.Domain, pair.Nameserver)
		}
		time.Sleep(time.Duration(delay) * time.Millisecond)
	}
	tracker <- 1337
}

func QueryRA(q <-chan Pair, tracker chan<- interface{}, delay int) {
	for pair := range q {
		msg := message(pair.Domain, dns.TypeA, true)
		for x := 0; x < 3; x++ {
			in, err := dns.Exchange(msg, pair.Nameserver+":53")
			if err != nil {
				Error("hiccup on " + pair.Nameserver + " retrying...")
				time.Sleep(1 * time.Second)
				continue
			}

			if len(in.Answer) > 0 {
				Found = true
				if in.Answer[0].Header().Ttl != Domains[pair.Domain].TTL {
					fmt.Printf("[%s] associated domain %s found on %s with mismatched TTL of %d\n", Vendors[Domains[pair.Domain].Vendor], pair.Domain, pair.Nameserver, in.Answer[0].Header().Ttl)
				}
				break
			}
		}
		time.Sleep(time.Duration(delay) * time.Millisecond)
	}
	tracker <- 1337
}

func Run(ra bool, threads, delay int) {
	pairs := make(chan Pair)
	tracker := make(chan interface{})

	if !ra {
		// non-recursive snoop
		Info(fmt.Sprintf("non-recursive snooping on %d resolvers...\n", len(Valid)))
		go func() {
			for i := 0; i < threads; i++ {
				Query(pairs, tracker, delay)
			}
		}()

		for _, ns := range Valid {
			for k, _ := range Domains {
				pairs <- Pair{Nameserver: ns, Domain: k}
			}
		}

		close(pairs)
	} else {
		Info(fmt.Sprintf("recursively snooping on %d resolvers...\n", len(Recursive)))
		go func() {
			for i := 0; i < threads; i++ {
				QueryRA(pairs, tracker, delay)
			}
		}()

		for _, ns := range Recursive {
			for k, _ := range Domains {
				pairs <- Pair{Nameserver: ns, Domain: k}
			}
		}

		close(pairs)
	}

	for x := 0; x < threads; x++ {
		<-tracker
	}
}