serverlist.go

package dnsyo

import (
	"fmt"
	"github.com/gocarina/gocsv"
	log "github.com/sirupsen/logrus"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

const (
	answerResult         = 4    // tab index for the result in a DNS response. Used to eliminate type and TTL etc.
	reliabilityThreshold = 0.97 // minimum public-dns.info reliability threshold for a server to be loaded from csv
)

// representation of a nameserver in CSV from from public-dns.info
type csvNameserver struct {
	// IPAddress is the ipv4 address of the server
	IPAddress string `csv:"ip"`

	// Name is the hostname of the server if the server has a hostname
	Name string `csv:"name"`

	// Country is the two-letter ISO 3166-1 alpha-2 code of the country
	Country string `csv:"country_id"`

	// City specifies the city that the server is hosted on
	City string `csv:"city"`

	// Version is the software version of the dns daemon that the server is using
	Version string `csv:"version"`

	// Error is the error that the server returned. Probably will be empty if you use the valid nameserver dataset
	Error string `csv:"error"`

	// DNSSec is a boolean to indicate if the server supports DNSSec or not
	DNSSec bool `csv:"dnssec"`

	// Realiability is a normalized value - from 0.0 - 1.0 - to indicate how stable the server is
	Reliability float64 `csv:"reliability"`

	// CheckedAt is a timestamp to indicate the date that the server was last checked
	CheckedAt time.Time `csv:"checked_at"`

	// CreatedAt is a timestamp to indicate when the server was inserted in the database
	CreatedAt time.Time `csv:"created_at"`
}

// ServerList is an alias for a slice of Server objects used for performing bulk actions on multiple threads
type ServerList []Server

// ServersFromFile loads a ServerList from a YAML file. Will raise an error if the file cannot be opened or processed
func ServersFromFile(filename string) (sl ServerList, err error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}

	err = yaml.Unmarshal(data, &sl)
	if err != nil {
		return nil, err
	}

	return
}

// ServersFromCSVURL loads a ServerList from a CSV provided at a given URL.
// Designed for use with public-info.dns
func ServersFromCSVURL(url string) (sl ServerList, err error) {
	csvFile, err := http.Get(url)
	if err != nil {
		return
	}

	// sometimes the file can be truncated and have an incomplete final line.
	// run a basic check to ensure it isn't going to error later.
	bytes, err := ioutil.ReadAll(csvFile.Body)
	if err != nil {
		return
	}

	rows := strings.Split(string(bytes), "\n")
	nCols := strings.Count(rows[0], ",")
	if strings.Count(rows[len(rows)-1], ",") < nCols {
		rows = rows[:len(rows)-2]
	}
	data := strings.Join(rows, "\n")

	var servers []csvNameserver
	err = gocsv.Unmarshal(strings.NewReader(data), &servers)
	if err != nil {
		return
	}

	for _, ns := range servers {
		if ip := net.ParseIP(ns.IPAddress); ip.To4() == nil {
			continue // we can't process IPv6 yet
		}
		if ns.Reliability >= reliabilityThreshold {
			s := Server{
				IP:      ns.IPAddress,
				Country: strings.ToUpper(ns.Country),
				Name:    ns.Name,
			}
			sl = append(sl, s)
		}
	}

	return
}

// DumpToFile a the current server list to a YAML file.
// Includes a commented header to identify the fact it is generated.
func (sl *ServerList) DumpToFile(filename string) (err error) {
	yml, err := yaml.Marshal(sl)
	if err != nil {
		return
	}

	data := append([]byte("#### GENERATED BY dnsyo update ####\n\n"), yml...)

	err = ioutil.WriteFile(filename, data, 0744)
	return
}

// FilterCountry filters the current server list by country and returns a new server list with the matching servers in it.
// Returns an error if no servers were found.
func (sl *ServerList) FilterCountry(country string) (fl ServerList, err error) {
	for _, s := range *sl {
		if s.Country == strings.ToUpper(country) {
			fl = append(fl, s)
		}
	}

	if len(fl) == 0 {
		err = fmt.Errorf("no servers matching country %s were found", country)
	}

	return
}

// NRandom returns n random servers from the current server list in a new list.
// Will return an error if there are less than n servers in the current list.
func (sl *ServerList) NRandom(n int) (rl ServerList, err error) {
	ql := *sl

	if len(ql) < n {
		return nil, fmt.Errorf("insufficient servers to populate list: %d of %d", len(ql), n)
	}

	rl = make(ServerList, n)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	perm := r.Perm(n)
	for i, randIndex := range perm[:n] {
		rl[i] = ql[randIndex]
	}

	return
}

// ExecuteQuery runs a Query object in a specified number of threads.
// The returned QueryResult is not associated with the provided Query, however may be set by the caller.
func (sl *ServerList) ExecuteQuery(q *Query, threads int) (qr QueryResults) {
	qr = make(QueryResults)

	var wg sync.WaitGroup
	var mtx sync.Mutex

	queue := make(chan Server, len(*sl))

	// start workers
	for i := 0; i < threads; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			for s := range queue {
				res, err := s.Lookup(q.Domain, q.Type)
				ans := strings.Join(res, "\n")

				r := new(Result)

				if err != nil {
					r.Error = err.Error()
				} else {
					r.Answer = ans
				}

				mtx.Lock()
				qr[s.String()] = r
				mtx.Unlock()
			}

		}(i)
	}

	for _, s := range *sl {
		queue <- s
	}
	close(queue)

	wg.Wait()
	return
}

//func (sl *ServerList) StreamQuery(q *Query, threads int, results chan QueryResults) {
//	recordType := dns.StringToType[q.Type]
//
//	var wg sync.WaitGroup
//
//	queue := make(chan Server, len(*sl))
//
//	// start workers
//	for i := 0; i < threads; i++ {
//		wg.Add(1)
//		go func(i int) {
//			defer wg.Done()
//			for s := range queue {
//				res, err := s.Lookup(q.Domain, recordType)
//				ans := strings.Join(res, "\n")
//
//				r := &Result{
//					Answer: ans,
//				}
//				if err != nil {
//					r.Error = err.Error()
//				}
//
//				qr := QueryResults{
//					s: r,
//				}
//
//				results <- qr
//			}
//		}(i)
//	}
//
//	for _, s := range *sl {
//		queue <- s
//	}
//	close(queue)
//
//	wg.Wait()
//	return
//}

// TestAll tests all the servers in the current list and returns a new list with only the workings ones.
func (sl *ServerList) TestAll(threads int) (working ServerList) {
	var wg sync.WaitGroup
	var mutex sync.Mutex
	testQueue := make(chan Server, len(*sl))

	// start workers
	for i := 0; i < threads; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			for s := range testQueue {
				log.WithField("thread", i).Debug("Testing " + s.Name)
				if ok, err := s.Test(); ok {
					mutex.Lock()
					working = append(working, s)
					mutex.Unlock()
				} else {
					log.WithFields(log.Fields{
						"thread": i,
						"server": s.String(),
						"reason": err,
					}).Info("Disabling server")
				}
			}
		}(i)
	}

	// add test servers
	for _, s := range *sl {
		testQueue <- s
	}
	close(testQueue)

	wg.Wait()
	return working
}