Add flag for custom HTTP headers

README.md

@@ -71,6 +71,8 @@ exit status 2
       Label to print (Optional) (default "none")
   -f string
       Format message with go template (Optional)
+  -h value
+      Custom headers to add, like 'X-Scope-OrgID: myorg' (Optional)
   -m string
       Comparison method (Optional) (default "ge")
   -u string
@@ -179,6 +181,15 @@ OK kube-node-2: CorruptDockerOverlay2 is OK
 OK kube-node-2: FrequentContainerdRestart is OK
 ```
 
+## Custom headers
+Prometheus-compatible systems may require extra headers, e.g. grafana's mimir when using tenants.
+
+Example:
+```
+-h "X-Scope-OrgID: myorg"
+```
+This will add the given header to the HTTP request.
+
 ## Method
 
 `-m ge OR gt OR le OR lt` tells the check how to compare the result with the critical and warning flags 

nagitheus.go

@@ -98,9 +98,22 @@ type TemplateData struct {
 	Critical string
 }
 
+type headers []string
+
+func (custom_headers *headers) String() string {
+	return strings.Join(*custom_headers, ", ")
+}
+
+func (custom_headers *headers) Set(value string) error {
+	*custom_headers = append(*custom_headers, value)
+	return nil
+}
+
 func main() {
+	var custom_headers headers
 
 	host := flag.String("H", "", "Host to query (Required, i.e. https://example.prometheus.com)")
+	flag.Var(&custom_headers, "h", "Custom headers to add, like 'X-Scope-OrgID: myorg' (Optional)")
 	query := flag.String("q", "", "Prometheus query (Required)")
 	warning := flag.String("w", "", "Warning treshold (Required)")
 	critical := flag.String("c", "", "Critical treshold (Required)")
@@ -129,7 +142,7 @@ func main() {
 	// check flags
 	flag.VisitAll(check_set)
 	// query prometheus
-	response := execute_query(*host, *query, *username, *password, *token)
+	response := execute_query(*host, custom_headers, *query, *username, *password, *token)
 	// print response (DEBUGGING)
 	if *debug == "yes" {
 		print_response(response)
@@ -155,15 +168,15 @@ func main() {
 }
 
 func check_set(argument *flag.Flag) {
-	if argument.Value.String() == "" && argument.Name != "u" && argument.Name != "p" && argument.Name != "t" &&
+	if argument.Value.String() == "" && argument.Name != "u" && argument.Name != "p" && argument.Name != "t" && argument.Name != "h" &&
 		argument.Name != "value-mapping" && argument.Name != "value-unit" && argument.Name != "max-chars" && argument.Name != "f" {
 		Message := "Please set value for : " + argument.Name
 		Usage()
 		exit_func(UNKNOWN, Message, 0)
 	}
 }
 
-func execute_query(host string, query string, username string, password string, token string) []byte {
+func execute_query(host string, custom_headers []string, query string, username string, password string, token string) []byte {
 	query_encoded := url.QueryEscape(query)
 	url_complete := host + "/api/v1/query?query=" + "(" + query_encoded + ")"
 
@@ -177,6 +190,12 @@ func execute_query(host string, query string, username string, password string,
 	}
 
 	req, err := http.NewRequest("GET", url_complete, nil)
+
+	for _, header := range custom_headers {
+		splitted := strings.Split(header, ":")
+		req.Header.Add(strings.TrimSpace(splitted[0]), strings.TrimSpace(splitted[1]))
+	}
+
 	if username != "" && password != "" {
 		req.SetBasicAuth(username, password)
 	}