From bd27c80b5d714e5f0bc0acff5d4d860efd02a0d1 Mon Sep 17 00:00:00 2001
From: Daniel Czerwonk <daniel@dan-nrw.de>
Date: Fri, 1 Dec 2017 07:42:29 +0100
Subject: [PATCH] added flags to disable protocols

---
 main.go              | 34 +++++++++++++++++++++++++++++++++-
 metric_collector.go  | 11 ++++++-----
 protocol/protocol.go |  8 ++++----
 3 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/main.go b/main.go
index ed79955..bed2a7b 100644
--- a/main.go
+++ b/main.go
@@ -9,6 +9,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/prometheus/common/log"
+	"github.com/czerwonk/bird_exporter/protocol"
 )
 
 const version string = "1.0.0"
@@ -22,6 +23,12 @@ var (
 	birdEnabled   = flag.Bool("bird.ipv4", true, "Get protocols from bird")
 	bird6Enabled  = flag.Bool("bird.ipv6", true, "Get protocols from bird6")
 	newFormat     = flag.Bool("use-new-format", false, "New metric format (more convinient / generic)")
+	enableBgp = flag.Bool("enable-bgp", true, "Enables metrics for protocol BGP")
+	enableOspf = flag.Bool("enable-ospf", true, "Enables metrics for protocol OSPF")
+	enableKernel = flag.Bool("enable-kernel", true, "Enables metrics for protocol kernel")
+	enableStatic = flag.Bool("enable-static", true, "Enables metrics for protocol static")
+	enableDevice = flag.Bool("enable-device", true, "Enables metrics for protocol static")
+	enableDirect = flag.Bool("enable-direct", true, "Enables metrics for protocol direct")
 )
 
 func init() {
@@ -88,7 +95,8 @@ func handleMetricsRequest(w http.ResponseWriter, r *http.Request) error {
 
 	if len(protocols) > 0 {
 		reg := prometheus.NewRegistry()
-		c := NewMetricCollectorForProtocols(protocols, *newFormat)
+		p := enabledProtocols()
+		c := NewMetricCollectorForProtocols(protocols, *newFormat, p)
 		reg.MustRegister(c)
 
 		promhttp.HandlerFor(reg, promhttp.HandlerOpts{
@@ -98,3 +106,27 @@ func handleMetricsRequest(w http.ResponseWriter, r *http.Request) error {
 
 	return nil
 }
+func enabledProtocols() int {
+	res := 0
+
+	if *enableBgp {
+		res |= protocol.BGP
+	}
+	if *enableOspf {
+		res |= protocol.OSPF
+	}
+	if *enableKernel {
+		res |= protocol.Kernel
+	}
+	if *enableStatic {
+		res |= protocol.Static
+	}
+	if *enableDevice {
+		res |= protocol.Device
+	}
+	if *enableDirect {
+		res |= protocol.Direct
+	}
+
+	return res
+}
diff --git a/metric_collector.go b/metric_collector.go
index 8712afb..159c053 100644
--- a/metric_collector.go
+++ b/metric_collector.go
@@ -8,11 +8,12 @@ import (
 )
 
 type MetricCollector struct {
-	protocols []*protocol.Protocol
-	exporters map[int][]metrics.MetricExporter
+	protocols        []*protocol.Protocol
+	exporters        map[int][]metrics.MetricExporter
+	enabledProtocols int
 }
 
-func NewMetricCollectorForProtocols(protocols []*protocol.Protocol, newFormat bool) *MetricCollector {
+func NewMetricCollectorForProtocols(protocols []*protocol.Protocol, newFormat bool, enabledProtocols int) *MetricCollector {
 	var e map[int][]metrics.MetricExporter
 
 	if newFormat {
@@ -21,7 +22,7 @@ func NewMetricCollectorForProtocols(protocols []*protocol.Protocol, newFormat bo
 		e = exportersForLegacy()
 	}
 
-	return &MetricCollector{protocols: protocols, exporters: e}
+	return &MetricCollector{protocols: protocols, exporters: e, enabledProtocols: enabledProtocols}
 }
 
 func exportersForLegacy() map[int][]metrics.MetricExporter {
@@ -61,7 +62,7 @@ func (m *MetricCollector) Describe(ch chan<- *prometheus.Desc) {
 
 func (m *MetricCollector) Collect(ch chan<- prometheus.Metric) {
 	for _, p := range m.protocols {
-		if p.Proto == protocol.PROTO_UNKNOWN {
+		if p.Proto == protocol.PROTO_UNKNOWN || (m.enabledProtocols & p.Proto != p.Proto) {
 			continue
 		}
 
diff --git a/protocol/protocol.go b/protocol/protocol.go
index a84e28f..2e309a4 100644
--- a/protocol/protocol.go
+++ b/protocol/protocol.go
@@ -4,10 +4,10 @@ const (
 	PROTO_UNKNOWN = 0
 	BGP           = 1
 	OSPF          = 2
-	Kernel        = 3
-	Static        = 4
-	Direct        = 5
-	Device        = 6
+	Kernel        = 4
+	Static        = 8
+	Direct        = 16
+	Device        = 32
 )
 
 type Protocol struct {