From 2e262a63a86f8bf4ede2a627841889662d2fc709 Mon Sep 17 00:00:00 2001
From: dave <dave@davepedu.com>
Date: Thu, 7 Mar 2024 22:34:51 -0700
Subject: [PATCH] add nagios mode

---
 cmd/etcdmtr/main.go | 45 +++++++++++++++++++++++++++++++++++----------
 1 file changed, 35 insertions(+), 10 deletions(-)

diff --git a/cmd/etcdmtr/main.go b/cmd/etcdmtr/main.go
index d06ef53..86939ae 100644
--- a/cmd/etcdmtr/main.go
+++ b/cmd/etcdmtr/main.go
@@ -25,6 +25,7 @@ type EtcdMtrArgs struct {
 	Timeout         int    `short:"t" long:"timeout" description:"Operation timeout in seconds"`
 	ExpectedMembers int    `short:"m" long:"members" description:"Expect at least this many members"`
 	Debug           bool   `long:"debug" description:"Show debug output and etcd messages"`
+	Nagios          bool   `long:"nagios" description:"Output nagios/nrpe format"`
 }
 
 type Client struct {
@@ -107,17 +108,33 @@ func main() {
 
 	logger := etcdmtr.NewRootLogger()
 
+	if args.Nagios {
+		fatalHandler := func(logger *etcdmtr.ELog, exitCode int, message string, a ...any) {
+			code2name := []string{"OK", "WARNING", "CRITICAL", "UNKNOWN"}
+			msg := fmt.Sprintf(message, a...)
+			fmt.Printf("%s - %s\n", code2name[exitCode], msg)
+			os.Exit(exitCode)
+		}
+		logger.FatalHandler = fatalHandler
+
+		logger.Level = etcdmtr.LOGLEVEL_SILENCE
+
+		if args.ExpectedMembers < 1 {
+			logger.Fatalc(3, "--nagios requires --members > 0")
+		}
+	}
+
 	if args.Version {
-		logger.Fatalc(1, "0.0.1")
+		logger.Fatalc(0, "0.0.1")
 	}
 
 	if len(positionals) < 1 {
-		logger.Fatal("must specify endpoints")
+		logger.Fatalc(3, "must specify endpoints")
 	}
 
 	if !args.Insecure {
 		if args.ClientKey == "" || args.ClientCert == "" || args.ServerCA == "" {
-			logger.Fatal("must set --insecure if not passing --client-cert, --client-key, and --ca")
+			logger.Fatalc(3, "must set --insecure if not passing --client-cert, --client-key, and --ca")
 		}
 	}
 
@@ -127,7 +144,7 @@ func main() {
 
 	client, err := NewClient(positionals, args.ClientCert, args.ClientKey, args.ServerCA, args.Debug)
 	if err != nil {
-		logger.Fatal("couldn't get client: %s", err.Error())
+		logger.Fatalc(3, "couldn't get client: %s", err.Error())
 	}
 
 	etcd, err := client.GetEtcd()
@@ -160,23 +177,31 @@ func main() {
 
 		is_leader, err := contactMember(contextCtx, client, endpoint)
 		if err != nil {
-			logger.Warning("couldn't reach member %s: %s", endpoint, err)
+			logger.Error("couldn't reach member %s: %s", endpoint, err)
 			memberErrors = append(memberErrors, err)
 			continue
 		}
 		logger.Info("reached member %s - is leader? %v", endpoint, is_leader)
 	}
 
+	if args.Nagios {
+		if numMembers != args.ExpectedMembers {
+
+		}
+	}
+
 	if args.ExpectedMembers > 0 {
 		if numMembers != args.ExpectedMembers {
-			logger.Error("found different number of members than expected: %d", numMembers)
-			os.Exit(1)
+			logger.Fatalc(2, "found %d members, expected %d", numMembers, args.ExpectedMembers)
 		}
 		if len(memberErrors) > 0 {
-			logger.Error("encountered an error when contacting %d members", len(memberErrors))
-			os.Exit(1)
+			logger.Fatalc(2, "couldn't contact %d members", len(memberErrors))
 		}
 	}
+
+	if args.Nagios {
+		logger.Fatalc(0, "found %d members", numMembers)
+	}
 }
 
 func contactMember(ctx context.Context, client *Client, endpoint string) (bool, error) {
@@ -188,7 +213,7 @@ func contactMember(ctx context.Context, client *Client, endpoint string) (bool,
 
 	status, err := etcd.Status(ctx, endpoint)
 	if err != nil {
-		return false, fmt.Errorf("couldn't reach member: %s", err)
+		return false, err
 	}
 	return status.Leader == status.Header.MemberId, nil
 }