//////////////////////////////////////////////////////////////////////////
// DNS Metrics
//////////////////////////////////////////////////////////////////////////

package main

//////////////////////////////////////////////////////////////////////////

import (
	"encoding/json"
	"errors"
	//	"fmt"
	dns "github.com/miekg/dns"
	"github.com/prometheus/client_golang/prometheus"
	log "github.com/sirupsen/logrus"
	"io/ioutil"
	"math"
	"net/http"
	"strconv"
	"strings"
	"time"
)

//////////////////////////////////////////////////////////////////////////
// data structures

// for holding the DNS metrics

type DNSMetrics struct {
	soa       *prometheus.GaugeVec
	rtt       *prometheus.GaugeVec
	valid     *prometheus.GaugeVec
	stime     *prometheus.GaugeVec
	label_map []prometheus.Labels
}

// structure for specifying DNS servers

type DNSServer struct {
	role  string
	owner string
	name  string
	ip    uint8
	addr  string
	soa   uint32
}

// structure for returning relevant DNS data

type DNSResult struct {
	serial uint32
	rtt    float64
	nsid   string
}

// data structures for querying the current commit metric

type DNSCommitMetrics struct {
	match        *prometheus.GaugeVec
	lastUpdate   time.Time
	mismatch     bool
	since        time.Time
	labelInvalid prometheus.Labels
	labelUpdate  prometheus.Labels
}

type DNSExplorerCommit struct {
	Commit string
}

type DNSMasterCommit struct {
	Type   string `json:"type"`
	URI    string `json:"uri"`
	Branch string `json:"branch"`
	Commit string `json:"commit"`
}

//////////////////////////////////////////////////////////////////////////

// hardcoded :( list of DNS servers to query

var dns_servers = []*DNSServer{

	// master

	&DNSServer{"master", "jrb0001", "b.master.delegation-servers.dn42", 6,
		"[fd42:180:3de0:30::1]:53", 0},
	&DNSServer{"master", "jrb0001", "j.master.delegation-servers.dn42", 6,
		"[fd42:180:3de0:10:5054:ff:fe87:ea39]:53", 0},

	// anycast

	&DNSServer{"recursive", "yamakaja", "a.recursive-servers.dn42", 4,
		"172.20.0.53:53", 0},
	&DNSServer{"recursive", "yamakaja", "a.recursive-servers.dn42", 6,
		"[fd42:d42:d42:54::1]:53", 0},

	// burble

	&DNSServer{"delegation", "burble", "b.delegation-servers.dn42", 4,
		"172.20.129.1:53", 0},
	&DNSServer{"delegation", "burble", "b.delegation-servers.dn42", 6,
		"[fd42:4242:2601:ac53::1]:53", 0},
	&DNSServer{"recursive", "burble", "b.recursive-servers.dn42", 4,
		"172.20.129.2:53", 0},
	&DNSServer{"recursive", "burble", "b.recursive-servers.dn42", 6,
		"[fd42:4242:2601:ac53::53]:53", 0},

	&DNSServer{"burble.dn42", "burble", "fr-rbx1", 6, "[fd42:4242:2601:36::ac:53]:53", 0},
	&DNSServer{"burble.dn42", "burble", "us-dal3", 6, "[fd42:4242:2601:2a::ac:53]:53", 0},
	&DNSServer{"burble.dn42", "burble", "sg-sin2", 6, "[fd42:4242:2601:37::ac:53]:53", 0},
	&DNSServer{"burble.dn42", "burble", "ca-bhs2", 6, "[fd42:4242:2601:2d::ac:53]:53", 0},
	&DNSServer{"burble.dn42", "burble", "lt-vil1", 6, "[fd42:4242:2601:3d::ac:53]:53", 0},

	// jrb0001

	&DNSServer{"delegation", "jrb0001", "j.delegation-servers.dn42", 4,
		"172.20.1.254:53", 0},
	&DNSServer{"delegation", "jrb0001", "j.delegation-servers.dn42", 6,
		"[fd42:5d71:219:1:a526:d935:281e:22d6]:53", 0},
	&DNSServer{"recursive", "jrb0001", "j.recursive-servers.dn42", 4,
		"172.20.1.255:53", 0},
	&DNSServer{"recursive", "jrb0001", "j.recursive-servers.dn42", 6,
		"[fd42:5d71:219:1:69c2:2b0e:17e8:c215]:53", 0},

	&DNSServer{"jrb0001", "jrb0001", "nl-1", 6, "[fd42:5d71:219::1]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "de-1", 6, "[fd42:5d71:219::2]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "gb-1", 6, "[fd42:5d71:219::3]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "fr-1", 6, "[fd42:5d71:219::4]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "us-1", 6, "[fd42:5d71:219::5]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "au-1", 6, "[fd42:5d71:219::6]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "jp-1", 6, "[fd42:5d71:219::7]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "sg-1", 6, "[fd42:5d71:219::8]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "ca-1", 6, "[fd42:5d71:219::9]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "us-2", 6, "[fd42:5d71:219::a]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "us-3", 6, "[fd42:5d71:219::b]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "us-4", 6, "[fd42:5d71:219::c]:53", 0},
	&DNSServer{"jrb0001", "jrb0001", "us-5", 6, "[fd42:5d71:219::d]:53", 0},

	// yamakaja

	&DNSServer{"delegation", "yamakaja", "y.delegation-servers.dn42", 4,
		"172.20.20.66:53", 0},
	&DNSServer{"delegation", "yamakaja", "y.delegation-servers.dn42", 6,
		"[fd42:c01d:beef::3]:53", 0},
	&DNSServer{"recursive", "yamakaja", "y.recursive-servers.dn42", 4,
		"172.20.20.65:53", 0},
	&DNSServer{"recursive", "yamakaja", "y.recursive-servers.dn42", 6,
		"[fd42:c01d:beef::2]:53", 0},
}

//////////////////////////////////////////////////////////////////////////
//
// DNS Server Metrics
//
//////////////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////////////
// initialisation function to register metrics

func (m *DNSMetrics) Register() {

	m.soa = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "dn42_dns_soa",
		Help: "SOA for .dn42 domain",
	}, []string{"role", "owner", "name", "ip", "addr"})
	prometheus.MustRegister(m.soa)

	m.rtt = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "dn42_dns_rtt",
		Help: "RTT when collecting SOA for .dn42 domain",
	}, []string{"role", "owner", "name", "ip", "addr"})
	prometheus.MustRegister(m.rtt)

	m.valid = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "dn42_dns_valid",
		Help: "0 = response and correct serial, 1 = response but incorrect serial, 2 = no response",
	}, []string{"role", "owner", "name", "ip", "addr"})
	prometheus.MustRegister(m.valid)

	m.stime = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "dn42_dns_stime",
		Help: "Returns the time between now and the SOA serial number",
	}, []string{"role", "owner", "name", "ip", "addr"})
	prometheus.MustRegister(m.stime)

	// pre-populate the labels
	m.label_map = make([]prometheus.Labels, len(dns_servers))

	for ix, server := range dns_servers {
		m.label_map[ix] = prometheus.Labels{
			"role":  server.role,
			"owner": server.owner,
			"name":  server.name,
			"ip":    strconv.Itoa(int(server.ip)),
			"addr":  server.addr,
		}
	}

}

//////////////////////////////////////////////////////////////////////////
// collect metrics for all DNS servers

func (m *DNSMetrics) Collect() {

	now := uint64(time.Now().Unix())

	// search the masters for the highest SOA
	var latest_soa uint32 = 0

	for _, server := range dns_servers {
		if server.role == "master" {
			if server.soa > latest_soa {
				latest_soa = server.soa
			}
		}
	}

	// hold the results in an array where each entry corresponds
	// to the dns_servers array
	results := make([]*DNSResult, len(dns_servers))

	// query each server up to 3 times to try and get a result
	for count := 0; count < 3; count++ {

		for ix := 0; ix < len(results); ix++ {
			if results[ix] == nil {
				// no result yet, query the server
				results[ix] = dns_servers[ix].Query()
			}
		}

	}

	// now go through each result and update the metrics
	for ix, r := range results {

		var valid uint = 2

		// check if there was a valid result
		if r != nil {

			// update the server SOA
			dns_servers[ix].soa = r.serial

			// SOA and RTT are direct metrics returned from the query
			m.soa.With(m.label_map[ix]).Set(float64(r.serial))
			m.rtt.With(m.label_map[ix]).Set(r.rtt)

			// if the server has a high rtt ( > 500ms), then log an info message
			if r.rtt > 500 {
				log.WithFields(log.Fields{
					"result": r,
					"server": dns_servers[ix].name,
					"ipv":    dns_servers[ix].ip,
				}).Info("DNS Server high RTT")
			}

			// check if the SOA matches the lastest master SOA
			if r.serial == latest_soa {
				valid = 0
			} else {
				valid = 1
			}

			// before setting whether the server is valid, calculate the stime
			// (difference in time between now and the SOA, to allow checking that
			// it is not stale)
			// it's possible that the SOA could be in the future if there is a
			// clock mismatch between monitor and DNS server, in which case this
			// is flagged as a server error

			if uint64(r.serial) > now {
				// server error
				valid = 2
			} else {
				m.stime.With(m.label_map[ix]).Set(float64(now - uint64(r.serial)))
			}

		}

		// finally set the valid status
		m.valid.With(m.label_map[ix]).Set(float64(valid))

	}

}

//////////////////////////////////////////////////////////////////////////
// construct the DNS query and send to a server

func (s *DNSServer) Query() *DNSResult {

	// create a new recursive query
	msg := new(dns.Msg)
	msg.Id = dns.Id()
	msg.RecursionDesired = (s.role == "recursive")

	// query the dn42 root zone SOA
	msg.Question = []dns.Question{
		{"dn42.", dns.TypeSOA, dns.ClassINET},
	}

	// add EDNS0 options to also query the service ID (NSID)
	// pretty much copied verbatim from the library docs
	//	opts := new(dns.OPT)
	//	opts.Hdr.Name = "."
	//	opts.Hdr.Rrtype = dns.TypeOPT

	// create the NSID option
	//	ns_opt := new(dns.EDNS0_NSID)
	//	ns_opt.Code = dns.EDNS0NSID
	//	ns_opt.Nsid = ""

	// add the NSID option to the opts RR
	//	opts.Option = append(opts.Option, ns_opt)

	// then add the opts RR to the query
	//	msg.SetEdns0(4096, false)
	//	msg.Extra = []dns.RR{opts}

	// create a new DNS client
	client := new(dns.Client)
	client.Timeout, _ = time.ParseDuration("4s")

	// and finally query the server
	resp, rtt, err := client.Exchange(msg, s.addr)
	if err != nil || len(resp.Answer) != 1 {
		log.WithFields(log.Fields{
			"error":  err,
			"resp":   resp,
			"server": s,
		}).Warn("Failed to query DNS server")

		return nil
	}

	//	fmt.Printf("Resp: %v\n", resp)

	// was an SOA returned ?
	if soa, ok := resp.Answer[0].(*dns.SOA); !ok {

		log.WithFields(log.Fields{
			"resp":   resp,
			"server": s,
		}).Warn("DNS response was not an SOA")

		return nil

	} else {
		// got an SOA result

		result := &DNSResult{
			serial: soa.Serial,
			rtt:    math.Round(rtt.Seconds() * 1000),
		}

		// did we also get an NSID result ?
		if opts := resp.IsEdns0(); opts != nil {
			// response contains an EDNS0 record

			// check for an NSID entry
			for _, s := range opts.Option {
				switch e := s.(type) {
				case *dns.EDNS0_NSID:
					result.nsid = e.Nsid
				}
			}
		}

		return result
	}
}

//////////////////////////////////////////////////////////////////////////
//
// DNS Commit Metric
//
// Check the master commit against the current registry commit
// (via the explorer) to check that the master is being updated
//
//////////////////////////////////////////////////////////////////////////

func (m *DNSCommitMetrics) Register() {
	m.match = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "dn42_dns_commit",
		Help: "Time since commits last matched (or -1 for error)",
	}, []string{"metric"})
	prometheus.MustRegister(m.match)

	m.labelInvalid = prometheus.Labels{
		"metric": "invalid",
	}

	m.labelUpdate = prometheus.Labels{
		"metric": "update",
	}

}

func (m *DNSCommitMetrics) Collect() {

	now := time.Now()
	interval := now.Sub(m.lastUpdate)

	// only check if it's more than 60 mins since the last successful match
	if (interval.Hours() >= 1.0) || m.mismatch {

		// fetch the current commit from the explorer
		ec, err := m.FetchExplorerCommit()
		if err != nil {
			// couldn't fetch the explorer commit value, set the metric to bad
			m.match.With(m.labelUpdate).Set(-1.0)
			m.match.With(m.labelInvalid).Set(-1.0)
			m.mismatch = true
			return
		}

		mc, err := m.FetchMasterCommit()
		if err != nil {
			// couldn't fetch master commit value, set the metric to bad
			m.match.With(m.labelUpdate).Set(-2.0)
			m.match.With(m.labelInvalid).Set(-1.0)
			m.mismatch = true
			return
		}

		m.lastUpdate = now
		m.match.With(m.labelUpdate).Set(0.0)

		if ec == mc {

			// if the commits match, null the metric
			m.match.With(m.labelInvalid).Set(0.0)
			m.mismatch = false

		} else {

			// was this the first mismatch ?
			if !m.mismatch {
				m.since = now
				m.mismatch = true
			}

			since := now.Sub(m.since).Seconds()

			log.WithFields(log.Fields{
				"master":   mc,
				"explorer": ec,
				"since":    since,
			}).Warn("DNS Commit Mismatch")

			// set the metric to be the interval since last good match
			m.match.With(m.labelInvalid).Set(float64(since))
		}

	} else {

		// update time since last check
		m.match.With(m.labelUpdate).Set(float64(interval.Seconds()))
	}

}

//////////////////////////////////////////////////////////////////////////
// fetch the current commit from the explorer

func (m *DNSCommitMetrics) FetchExplorerCommit() (string, error) {

	response, err := http.Get("http://collector.dn42:8043/api/registry/.meta")
	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Warn("Unable to query registry explorer")
		return "", err
	}

	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Warn("Unable read explorer response")
		return "", err
	}

	var ec DNSExplorerCommit
	if err := json.Unmarshal(data, &ec); err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Warn("Unable to parse explorer JSON")
		return "", err
	}

	return strings.ToLower(ec.Commit), nil

}

//////////////////////////////////////////////////////////////////////////
// fetch the current commit from the master

func (m *DNSCommitMetrics) FetchMasterCommit() (string, error) {

	response, err :=
		http.Get("http://[fd42:180:3de0:10:5054:ff:fe87:ea39]:8080/api/git-db-state")
	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Warn("Unable to query registry master")
		return "", err
	}

	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Warn("Unable read master response")
		return "", err
	}

	var mc []DNSMasterCommit
	if err := json.Unmarshal(data, &mc); err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Warn("Unable to parse master JSON")
		return "", err
	}

	for _, c := range mc {
		if c.URI == "https://git.dn42.us/dn42/registry.git" {
			return strings.ToLower(c.Commit), nil
		}
	}

	log.WithFields(log.Fields{
		"MasterCommit": mc,
	}).Warn("Unable to find registry commit from master")
	return "", errors.New("Unable to find registry commit from master")
}

//////////////////////////////////////////////////////////////////////////
// end of code