1package main
2
3//go:generate go-bindata --prefix config/ config/
4
5import (
6	"bufio"
7	"flag"
8	"fmt"
9	"log"
10	"net/http"
11	"strconv"
12	"strings"
13
14	"github.com/optix2000/go-nsdctl"
15	"github.com/prometheus/client_golang/prometheus"
16	"github.com/prometheus/client_golang/prometheus/promhttp"
17)
18
19// Args
20var listenAddr = flag.String("listen-address", ":8080", "The address to listen on for HTTP requests.")
21var metricPath = flag.String("metric-path", "/metrics", "The path to export Prometheus metrocs to.")
22var metricConfigPath = flag.String("metric-config", "", "Mapping file for metrics. Defaults to built in file for NSD 4.1.x. This allows you to add or change any metrics that this scrapes")
23var nsdConfig = flag.String("config-file", "/usr/local/etc/nsd/nsd.conf", "Configuration file for nsd/unbound to autodetect configuration from. Defaults to /etc/nsd/nsd.conf. Mutually exclusive with -nsd-address, -cert, -key and -ca")
24var nsdType = flag.String("type", "nsd", "What nsd-like daemon to scrape (nsd or unbound). Defaults to nsd")
25var cert = flag.String("cert", "", "Client cert file location. Mutually exclusive with -config-file.")
26var key = flag.String("key", "", "Client key file location. Mutually exclusive with -config-file.")
27var ca = flag.String("ca", "", "Server CA file location. Mutually exclusive with -config-file.")
28var nsdAddr = flag.String("nsd-address", "", "NSD or Unbound control socket address.")
29
30// Prom stuff
31var nsdToProm = strings.NewReplacer(".", "_")
32
33var metricConfiguration = &metricConfig{}
34
35type NSDCollector struct {
36	client  *nsdctl.NSDClient
37	metrics map[string]*promMetric // Map of metric names to prom metrics
38}
39
40type promMetric struct {
41	desc      *prometheus.Desc
42	valueType prometheus.ValueType
43	labels    []string
44}
45
46func (c *NSDCollector) Describe(ch chan<- *prometheus.Desc) {
47	for _, metric := range c.metrics {
48		ch <- metric.desc
49	}
50}
51
52func (c *NSDCollector) Collect(ch chan<- prometheus.Metric) {
53	r, err := c.client.Command("stats_noreset")
54	if err != nil {
55		log.Println(err)
56		return
57	}
58
59	s := bufio.NewScanner(r)
60	for s.Scan() {
61		line := strings.Split(s.Text(), "=")
62		metricName := strings.TrimSpace(line[0])
63		m, ok := c.metrics[metricName]
64		if !ok {
65			log.Println("New metric " + metricName + " found. Refreshing.")
66			// Try to update the metrics list
67			err = c.updateMetric(s.Text())
68			if err != nil {
69				log.Println(err.Error())
70			}
71			// Refetch metric
72			m, ok = c.metrics[metricName]
73			if !ok {
74				log.Println("Metric " + metricName + "not configured. Skipping")
75			}
76			continue
77		}
78		value, err := strconv.ParseFloat(line[1], 64)
79		if err != nil {
80			log.Println(err)
81			continue
82		}
83		metric, err := prometheus.NewConstMetric(m.desc, m.valueType, value, m.labels...)
84		if err != nil {
85			log.Println(err)
86			continue
87		}
88		ch <- metric
89	}
90	err = s.Err()
91	if err != nil {
92		log.Println(err)
93		return
94	}
95
96}
97
98func (c *NSDCollector) updateMetric(s string) error {
99	// Assume line is in "metric=#" format
100	line := strings.Split(s, "=")
101	metricName := strings.TrimSpace(line[0])
102
103	_, exists := c.metrics[metricName]
104	if !exists {
105		metricConf, ok := metricConfiguration.Metrics[metricName]
106		if ok {
107			promName := nsdToProm.Replace(line[0])
108			c.metrics[metricName] = &promMetric{
109				desc: prometheus.NewDesc(
110					prometheus.BuildFQName(*nsdType, "", promName),
111					metricConf.Help,
112					nil,
113					nil,
114				),
115				valueType: metricConf.Type,
116			}
117		} else { // Try labeled metric
118			for _, v := range metricConfiguration.LabelMetrics {
119				labels := v.Regex.FindStringSubmatch(metricName)
120				if labels != nil {
121					var promName string
122					if v.Name != "" {
123						promName = v.Name
124					} else {
125						promName = nsdToProm.Replace(line[0])
126					}
127					c.metrics[metricName] = &promMetric{
128						desc: prometheus.NewDesc(
129							prometheus.BuildFQName(*nsdType, "", promName),
130							v.Help,
131							v.Labels,
132							nil,
133						),
134						valueType: v.Type,
135						labels:    labels[1:len(labels)],
136					}
137					// python "for-else"
138					goto Found
139				}
140			}
141			return fmt.Errorf("Metric ", metricName, " not found in config.")
142		Found:
143		}
144	}
145	return nil
146}
147
148func (c *NSDCollector) initMetricsList() error {
149	r, err := c.client.Command("stats_noreset")
150	if err != nil {
151		log.Println(err)
152		return err
153	}
154
155	if c.metrics == nil {
156		c.metrics = make(map[string]*promMetric)
157	}
158
159	// Grab metrics
160	s := bufio.NewScanner(r)
161	for s.Scan() {
162		err = c.updateMetric(s.Text())
163		if err != nil {
164			log.Println(err.Error(), "Skipping.")
165		}
166	}
167	return s.Err()
168}
169
170func NewNSDCollector(nsdType string, hostString string, caPath string, keyPath string, certPath string, skipVerify bool) (*NSDCollector, error) {
171	client, err := nsdctl.NewClient(nsdType, hostString, caPath, keyPath, certPath, skipVerify)
172	if err != nil {
173		return nil, err
174	}
175
176	collector := &NSDCollector{
177		client: client,
178	}
179
180	err = collector.initMetricsList()
181	if err != nil {
182		log.Println(err)
183		return nil, err
184	}
185	return collector, err
186}
187
188func NewNSDCollectorFromConfig(path string) (*NSDCollector, error) {
189	client, err := nsdctl.NewClientFromConfig(path)
190	if err != nil {
191		return nil, err
192	}
193
194	collector := &NSDCollector{
195		client: client,
196	}
197
198	err = collector.initMetricsList()
199	if err != nil {
200		log.Println(err)
201		return nil, err
202	}
203	return collector, err
204}
205
206// Main
207
208func main() {
209	flag.Parse()
210
211	// Load config
212	err := loadConfig(*metricConfigPath, metricConfiguration)
213	if err != nil {
214		log.Fatal(err)
215	}
216
217	// If one is set, all must be set.
218	var nsdCollector *NSDCollector
219	if *cert != "" || *key != "" || *ca != "" || *nsdAddr != "" {
220		if *cert != "" && *key != "" && *ca != "" && *nsdAddr != "" {
221			// Build from arguments
222			nsdCollector, err = NewNSDCollector(*nsdType, *nsdAddr, *ca, *key, *cert, false)
223			if err != nil {
224				log.Fatal(err)
225			}
226		} else {
227			log.Fatal("-cert, -key, and -ca must all be defined.")
228		}
229	} else {
230		// Build from config
231		nsdCollector, err = NewNSDCollectorFromConfig(*nsdConfig)
232		if err != nil {
233			log.Fatal(err)
234		}
235	}
236	prometheus.MustRegister(nsdCollector)
237	log.Println("Started.")
238	http.Handle(*metricPath, promhttp.Handler())
239	log.Fatal(http.ListenAndServe(*listenAddr, nil))
240}
241