package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
	"strings"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/miekg/dns"
	prom "github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/promlog"
	"github.com/prometheus/common/promlog/flag"
	"github.com/prometheus/common/version"
	"github.com/prometheus/exporter-toolkit/web"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
	"gopkg.in/yaml.v2"
)

const namespace = "dnssec"

type MetricsCollector struct {
	Log               log.Logger
	signatureOkMetric *prom.GaugeVec
	sokRcodeMetric    *prom.GaugeVec
	dsCountMetric     *prom.GaugeVec
	dsRcodeMetric     *prom.GaugeVec
	cdsCountMetric    *prom.GaugeVec
	cdsRcodeMetric    *prom.GaugeVec
	cdsDsMatchMetric  *prom.GaugeVec
	zones             []string
	DNS               dns.Client
	Resolver          string
}

type DnsConfig struct {
	Resolver string   `yaml:"resolver"`
	Zones    []string `yaml:"zones"`
	RRTypes  []string `yaml:"rrtypes"`
}

type DnssecConfig struct {
	DnsConfig DnsConfig `yaml:"dns"`
}

func b2f64(b bool) float64 {
	if b {
		return 1
	}
	return 0
}

func NewMetricsCollector(zonelist []string, resolver string, logger log.Logger) *MetricsCollector {
	signatureOkMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "signature_ok",
			Help:      "1 if the DNSSEC signature is present and valid, 0 otherwise",
		},
		[]string{"zone", "tld", "parent"},
	)
	sokRcodeMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "signature_rcode",
			Help:      "RCode of the DNS query",
		},
		[]string{"zone", "tld", "parent"},
	)
	dsCountMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "ds_count",
			Help:      "Number of DS record is present in the parent zone",
		},
		[]string{"zone", "tld", "parent"},
	)
	dsRcodeMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "ds_rcode",
			Help:      "RCode of the DS record answer",
		},
		[]string{"zone", "tld", "parent"},
	)
	cdsCountMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "cds_count",
			Help:      "Number of CDS records present in the zone",
		},
		[]string{"zone", "tld", "parent"},
	)
	cdsRcodeMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "cds_rcode",
			Help:      "RCode of the CDS record answer",
		},
		[]string{"zone", "tld", "parent"},
	)
	cdsDsMatchMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "cds_ds_match",
			Help:      "1 if the CDS and DS records match, 0 otherwise",
		},
		[]string{"zone", "tld", "parent"},
	)
	collector := MetricsCollector{
		signatureOkMetric: signatureOkMetric,
		sokRcodeMetric:    sokRcodeMetric,
		dsCountMetric:     dsCountMetric,
		dsRcodeMetric:     dsRcodeMetric,
		cdsCountMetric:    cdsCountMetric,
		cdsRcodeMetric:    cdsRcodeMetric,
		cdsDsMatchMetric:  cdsDsMatchMetric,
		zones:             zonelist,
		Resolver:          resolver,
		Log:               logger,
	}
	return &collector
}

func (c *MetricsCollector) Describe(ch chan<- *prom.Desc) {
	c.signatureOkMetric.Describe(ch)
	c.sokRcodeMetric.Describe(ch)
	c.dsCountMetric.Describe(ch)
	c.dsRcodeMetric.Describe(ch)
	c.cdsCountMetric.Describe(ch)
	c.cdsRcodeMetric.Describe(ch)
	c.cdsDsMatchMetric.Describe(ch)
}

func (c *MetricsCollector) Collect(ch chan<- prom.Metric) {
	for _, zone := range c.zones {
		labels := strings.Split(zone, ".")
		tld := strings.Join(labels[len(labels)-2:], ".")
		parent := strings.Join(labels[1:], ".")
		promlabels := prom.Labels{
			"zone":   zone,
			"tld":    tld,
			"parent": parent,
		}
		sok, sokrcode, _ := c.signatureOk(zone)
		c.sokRcodeMetric.With(promlabels).Set(float64(sokrcode))
		c.signatureOkMetric.With(promlabels).Set(b2f64(sok))
		dscount, dsrcode, cdscount, cdsrcode, cdm, _ := c.cdsDsMatches(zone)
		c.dsCountMetric.With(promlabels).Set(float64(dscount))
		c.dsRcodeMetric.With(promlabels).Set(float64(dsrcode))
		c.cdsCountMetric.With(promlabels).Set(float64(cdscount))
		c.cdsRcodeMetric.With(promlabels).Set(float64(cdsrcode))
		c.cdsDsMatchMetric.With(promlabels).Set(b2f64(cdm))
	}
	c.signatureOkMetric.Collect(ch)
	c.sokRcodeMetric.Collect(ch)
	c.dsCountMetric.Collect(ch)
	c.dsRcodeMetric.Collect(ch)
	c.cdsCountMetric.Collect(ch)
	c.cdsRcodeMetric.Collect(ch)
	c.cdsDsMatchMetric.Collect(ch)
}

type DsData struct {
	KeyTag     uint16
	Algorithm  uint8
	DigestType uint8
	Digest     string
}

func Equal(a, b []DsData) bool {
	if len(a) != len(b) {
		return false
	}
	for i, v := range a {
		if v != b[i] {
			return false
		}
	}
	return true
}

func (c *MetricsCollector) signatureOk(zone string) (bool, int, error) {
	qsoa := new(dns.Msg)
	// Request DNSSEC information
	qsoa.SetEdns0(4096, true)
	qsoa.SetQuestion(dns.Fqdn(zone), dns.TypeSOA)
	rsoa, _, err := c.DNS.Exchange(qsoa, c.Resolver)
	if err != nil {
		return false, rsoa.MsgHdr.Rcode, err
	}
	_ = level.Debug(c.Log).Log(fmt.Sprintf("SOA header: %+v\n", rsoa.MsgHdr))
	return rsoa.MsgHdr.AuthenticatedData && rsoa.MsgHdr.Rcode == dns.RcodeSuccess, rsoa.MsgHdr.Rcode, nil
}

func (c *MetricsCollector) cdsDsMatches(zone string) (int, int, int, int, bool, error) {
	dsCount := 0
	cdsCount := 0
	match := false
	var dsdata []DsData
	var cdsdata []DsData

	qcds := new(dns.Msg)
	qcds.SetQuestion(dns.Fqdn(zone), dns.TypeCDS)
	rcds, _, err1 := c.DNS.Exchange(qcds, c.Resolver)
	cdsRcode := rcds.MsgHdr.Rcode
	if err1 == nil {
		cdsCount = len(rcds.Answer)
		cdsdata = make([]DsData, len(rcds.Answer))
		_ = level.Debug(c.Log).Log(fmt.Sprintf("CDS data: %+v\n", rcds.Answer))
		for i, a := range rcds.Answer {
			ds := a.(*dns.CDS).DS
			cdsdata[i] = DsData{
				KeyTag:     ds.KeyTag,
				Algorithm:  ds.Algorithm,
				DigestType: ds.DigestType,
				Digest:     ds.Digest,
			}
		}
		sort.SliceStable(cdsdata, func(a, b int) bool { return cdsdata[a].KeyTag < cdsdata[b].KeyTag })
	}

	qds := new(dns.Msg)
	qds.SetQuestion(dns.Fqdn(zone), dns.TypeDS)
	rds, _, err2 := c.DNS.Exchange(qds, c.Resolver)
	dsRcode := rds.MsgHdr.Rcode
	if err2 == nil {
		dsCount = len(rds.Answer)
		dsdata = make([]DsData, len(rds.Answer))
		_ = level.Debug(c.Log).Log(fmt.Sprintf("DS data: %+v\n", rds.Answer))
		for i, a := range rds.Answer {
			ds := a.(*dns.DS)
			dsdata[i] = DsData{
				KeyTag:     ds.KeyTag,
				Algorithm:  ds.Algorithm,
				DigestType: ds.DigestType,
				Digest:     ds.Digest,
			}
		}
		sort.SliceStable(dsdata, func(a, b int) bool { return dsdata[a].KeyTag < dsdata[b].KeyTag })
	}

	if dsdata != nil && cdsdata != nil {
		match = cdsCount > 0 && Equal(dsdata, cdsdata)
		// special case: removal requested
		if len(cdsdata) == 1 && len(dsdata) == 0 && cdsdata[0].Algorithm == 0 {
			match = true
		}
	}
	return dsCount, dsRcode, cdsCount, cdsRcode, match, nil
}

func main() {
	var (
		listenAddress = kingpin.Flag(
			"web.listen-address",
			"Address on which to expose metrics.",
		).Default(":9142").String()
		metricsPath = kingpin.Flag(
			"web.metrics-path",
			"Path under which to expose metrics.",
		).Default("/metrics").String()
		webConfigFile = kingpin.Flag(
			"web.config",
			"Path to web config yaml file.",
		).Default("").String()
		configFile = kingpin.Flag(
			"config",
			"Path to config yaml file.",
		).Default("").String()
	)

	promlogConfig := &promlog.Config{}
	flag.AddFlags(kingpin.CommandLine, promlogConfig)
	kingpin.Version(version.Print("dnssec_exporter"))
	kingpin.CommandLine.UsageWriter(os.Stdout)
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	logger := promlog.New(promlogConfig)
	logger = level.NewFilter(logger, level.AllowDebug())

	config := DnssecConfig{
		DnsConfig: DnsConfig{
			Zones:    []string{"example.org"},
			Resolver: "1.1.1.1:53",
		},
	}
	if configFile != nil {
		content, err := ioutil.ReadFile(*configFile)
		if err != nil {
			_ = level.Error(logger).Log(err.Error())
			os.Exit(1)
		}
		err = yaml.Unmarshal(content, &config)
		if err != nil {
			_ = level.Error(logger).Log(err.Error())
			os.Exit(1)
		}
	}
	// Always add trailing dot
	for i, z := range config.DnsConfig.Zones {
		if z[len(z)-1] != byte('.') {
			config.DnsConfig.Zones[i] = z + "."
		}
	}

	metricsCollector := NewMetricsCollector(config.DnsConfig.Zones, config.DnsConfig.Resolver, logger)
	prom.MustRegister(metricsCollector)
	http.Handle(*metricsPath, promhttp.Handler())
	server := &http.Server{Addr: *listenAddress}
	_ = logger.Log(web.ListenAndServe(server, *webConfigFile, logger))
}