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

package main

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

import (
	//	"fmt"
	dns "github.com/miekg/dns"
	"github.com/prometheus/client_golang/prometheus"
	log "github.com/sirupsen/logrus"
	"math"
	"strconv"
)

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

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

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

var dns_servers = []*DNSServer{
	&DNSServer{"master", "master.delegation-servers.dn42", 6, "[fd42:180:3de0:10:5054:ff:fe87:ea39]:53", 0},
	&DNSServer{"delegation", "b.delegation-servers.dn42", 4, "172.20.129.1:53", 0},
	&DNSServer{"delegation", "b.delegation-servers.dn42", 6, "[fd42:4242:2601:ac53::1]:53", 0},
	&DNSServer{"delegation", "j.delegation-servers.dn42", 4, "172.20.1.18:53", 0},
	&DNSServer{"delegation", "j.delegation-servers.dn42", 6, "[fd42:5d71:219:0:1::42]:53", 0},
	&DNSServer{"delegation", "y.delegation-servers.dn42", 4, "172.20.20.66:53", 0},
	&DNSServer{"delegation", "y.delegation-servers.dn42", 6, "[fd42:c01d:beef::3]:53", 0},
	&DNSServer{"recursive", "a.recursive-servers.dn42", 4, "172.20.0.53:53", 0},
	&DNSServer{"recursive", "a.recursive-servers.dn42", 6, "[fd42:d42:d42:54::1]:53", 0},
	&DNSServer{"recursive", "b.recursive-servers.dn42", 4, "172.20.129.2:53", 0},
	&DNSServer{"recursive", "b.recursive-servers.dn42", 6, "[fd42:4242:2601:ac53::53]:53", 0},
	&DNSServer{"recursive", "j.recursive-servers.dn42", 4, "172.20.1.19:53", 0},
	&DNSServer{"recursive", "j.recursive-servers.dn42", 6, "[fd42:5d71:219:0:1::43]:53", 0},
	&DNSServer{"recursive", "y.recursive-servers.dn42", 4, "172.20.20.65:53", 0},
	&DNSServer{"recursive", "y.recursive-servers.dn42", 6, "[fd42:c01d:beef::2]:53", 0},
}

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

func (m *DNSMetrics) Register() {

	m.soa = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "dn42_dns_soa",
		Help: "SOA for .dn42 domain",
	}, []string{"role", "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", "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", "name", "ip", "addr"})
	prometheus.MustRegister(m.valid)

	// 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,
			"name": server.name,
			"ip":   strconv.Itoa(int(server.ip)),
			"addr": server.addr,
		}
	}

}

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

func (m *DNSMetrics) Collect() {

	for ix, server := range dns_servers {
		soa, rtt := server.Query()
		server.soa = uint(soa)

		m.soa.With(m.label_map[ix]).Set(soa)
		m.rtt.With(m.label_map[ix]).Set(rtt)

		var valid uint = 0

		if server.soa == 0 {
			valid = 2
		} else {
			if server.soa != dns_servers[0].soa {
				valid = 1
			} else {
				valid = 0
			}
		}

		m.valid.With(m.label_map[ix]).Set(float64(valid))
	}

}

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

func (s *DNSServer) Query() (float64, float64) {

	msg := new(dns.Msg)
	msg.Id = dns.Id()
	msg.RecursionDesired = (s.role == "recursive")
	msg.Question = []dns.Question{{"dn42.", dns.TypeSOA, dns.ClassINET}}

	client := new(dns.Client)
	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("Unable to query DNS server")

		return 0, 0
	}

	if soa, ok := resp.Answer[0].(*dns.SOA); ok {
		return float64(soa.Serial), math.Round(rtt.Seconds() * 1000)
	}

	return 0, 0
}

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