package main

import (
	"fmt"
	"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"
)

const namespace = "dnssec"

var removalCds = DsData{
	KeyTag:     0,
	Algorithm:  0,
	DigestType: 0,
	Digest:     "00",
}

type MetricsCollector struct {
	Log               log.Logger
	signatureOkMetric *prom.GaugeVec
	dsPresentMetric   *prom.GaugeVec
	cdsPresentMetric  *prom.GaugeVec
	cdsDsMatchMetric  *prom.GaugeVec
	zones             []string
	DNS               dns.Client
	Resolver          string
}

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"},
	)
	dsPresentMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "ds_present",
			Help:      "1 if  CDS record is present in the parent zone, 0 otherwise",
		},
		[]string{"zone", "tld", "parent"},
	)
	cdsPresentMetric := prom.NewGaugeVec(
		prom.GaugeOpts{
			Namespace: namespace,
			Name:      "cds_present",
			Help:      "1 if a CDS record is present in the zone, 0 otherwise",
		},
		[]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,
		dsPresentMetric:   dsPresentMetric,
		cdsPresentMetric:  cdsPresentMetric,
		cdsDsMatchMetric:  cdsDsMatchMetric,
		zones:             zonelist,
		Resolver:          resolver,
		Log:               logger,
	}
	return &collector
}

func (c *MetricsCollector) Describe(ch chan<- *prom.Desc) {
	c.signatureOkMetric.Describe(ch)
	c.dsPresentMetric.Describe(ch)
	c.cdsPresentMetric.Describe(ch)
	c.cdsDsMatchMetric.Describe(ch)
}

func (c *MetricsCollector) Collect(ch chan<- prom.Metric) {
	err := c.reportMetrics(ch)
	if err != nil {
		// TODO
		c.Log.Log(err)
	}
}

func (c *MetricsCollector) reportMetrics(ch chan<- prom.Metric) error {
	for _, zone := range c.zones {
		sok, err1 := c.signatureOk(zone)
		if err1 != nil {
			continue
		}
		dsp, cdp, cdm, err2 := c.cdsDsMatches(zone)
		if err2 != nil {
			continue
		}
		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,
		}
		signatureOkMetric := c.signatureOkMetric.With(promlabels)
		dsPresentMetric := c.dsPresentMetric.With(promlabels)
		cdsPresentMetric := c.cdsPresentMetric.With(promlabels)
		cdsDsMatchMetric := c.cdsDsMatchMetric.With(promlabels)
		if sok {
			signatureOkMetric.Set(1)
		} else {
			signatureOkMetric.Set(0)
		}
		if dsp {
			dsPresentMetric.Set(1)
		} else {
			dsPresentMetric.Set(0)
		}
		if cdp {
			cdsPresentMetric.Set(1)
		} else {
			cdsPresentMetric.Set(0)
		}
		if cdm {
			cdsDsMatchMetric.Set(1)
		} else {

			cdsDsMatchMetric.Set(0)
		}
	}
	c.signatureOkMetric.Collect(ch)
	c.dsPresentMetric.Collect(ch)
	c.cdsPresentMetric.Collect(ch)
	c.cdsDsMatchMetric.Collect(ch)
	return nil
}

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, 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, err
	}
	level.Debug(c.Log).Log(fmt.Sprintf("SOA header: %+v\n", rsoa.MsgHdr))
	return rsoa.MsgHdr.AuthenticatedData && rsoa.MsgHdr.Rcode == dns.RcodeSuccess, nil
}

func (c *MetricsCollector) cdsDsMatches(zone string) (bool, bool, bool, error) {
	dsPresent := false
	cdsPresent := false
	
	qcds := new(dns.Msg)
	qcds.SetQuestion(dns.Fqdn(zone), dns.TypeCDS)
	rcds, _, err := c.DNS.Exchange(qcds, c.Resolver)
	if err != nil {
		return dsPresent, cdsPresent, false, err
	}
	cdsPresent = len(rcds.Answer) > 0
	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, _, err := c.DNS.Exchange(qds, c.Resolver)
	if err != nil {
		return dsPresent, cdsPresent, false, err
	}
	dsPresent = len(rds.Answer) > 0
	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 })

	match := cdsPresent && Equal(dsdata, cdsdata)
	// special case: removal requested
	if len(cdsdata) == 1 && len(dsdata) == 0 && cdsdata[0] == removalCds {
		match = true
	}
	return dsPresent, cdsPresent, 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()
		resolver = kingpin.Flag(
			"dns.resolver",
			"DNS Resolver to use.",
		).Default("127.0.0.1:53").String()
		zones = kingpin.Flag(
			"dns.zones",
			"DNS zones to check, comma separated.",
		).Default("example.org").String()
		configFile = kingpin.Flag(
			"web.config",
			"[EXPERIMENTAL] Path to config yaml file that can enable TLS or authentication.",
		).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()
	zonelist := strings.Split(*zones, ",")
	// Always add trailing dot
	for i, z := range zonelist {
		if z[len(z)-1] != byte('.') {
			zonelist[i] = z + "."
		}
	}
	logger := promlog.New(promlogConfig)
	logger = level.NewFilter(logger, level.AllowDebug())
	logger = log.With(logger, "ts", log.DefaultTimestampUTC)
	metricsCollector := NewMetricsCollector(zonelist, *resolver, logger)
	prom.MustRegister(metricsCollector)
	http.Handle(*metricsPath, promhttp.Handler())
	server := &http.Server{Addr: *listenAddress}
	logger.Log(web.ListenAndServe(server, *configFile, logger))
}