From 55312a623d4d3f678c4338096c37c5257b8b0030 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Benkovsk=C3=BD?= <ondrej.benkovsky@jamf.com>
Date: Wed, 24 Apr 2024 18:09:28 +0200
Subject: [PATCH] introduce a fail flag controlling when dnspyre exits with
 non-zero code

---
 cmd/root.go             | 39 +++++++++++++++++++++++++++++++++++++++
 docs/failoncondition.md | 29 +++++++++++++++++++++++++++++
 docs/index.md           |  4 ++++
 3 files changed, 72 insertions(+)
 create mode 100644 docs/failoncondition.md

diff --git a/cmd/root.go b/cmd/root.go
index e98cd72..8c79cd0 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -28,6 +28,15 @@ var (
 	benchmark = dnsbench.Benchmark{
 		Writer: os.Stdout,
 	}
+
+	failConditions []string
+)
+
+const (
+	ioerrorFailCondition    = "ioerror"
+	negativeFailCondition   = "negative"
+	errorFailCondition      = "error"
+	idmismatchFailCondition = "idmismatch"
 )
 
 func init() {
@@ -130,6 +139,12 @@ func init() {
 	pApp.Flag("progress", "Controls whether the progress bar is shown. Enabled by default.").
 		Default("true").BoolVar(&benchmark.ProgressBar)
 
+	pApp.Flag("fail", "Controls conditions upon which the dnspyre will exit with a non-zero exit code. Repeatable flag. "+
+		"Supported options are 'ioerror' (fail if there is at least 1 IO error), 'negative' (fail if there is at least 1 negative DNS answer), "+
+		"'error' (fail if there is at least 1 error DNS response), 'idmismatch' (fail there is at least 1 ID mismatch between DNS request and response).").
+		PlaceHolder(ioerrorFailCondition).
+		EnumsVar(&failConditions, ioerrorFailCondition, negativeFailCondition, errorFailCondition, idmismatchFailCondition)
+
 	pApp.Arg("queries", "Queries to issue. It can be a local file referenced using @<file-path>, for example @data/2-domains. "+
 		"It can also be resource accessible using HTTP, like https://raw.githubusercontent.com/Tantalor93/dnspyre/master/data/1000-domains, in that "+
 		"case, the file will be downloaded and saved in-memory. "+
@@ -183,6 +198,30 @@ func Execute() {
 	}
 
 	close(sigsInt)
+
+	if len(failConditions) > 0 {
+		stats := reporter.Merge(&benchmark, res)
+		for _, f := range failConditions {
+			switch f {
+			case ioerrorFailCondition:
+				if stats.Counters.IOError > 0 {
+					os.Exit(1)
+				}
+			case negativeFailCondition:
+				if stats.Counters.Negative > 0 {
+					os.Exit(1)
+				}
+			case errorFailCondition:
+				if stats.Counters.Error > 0 {
+					os.Exit(1)
+				}
+			case idmismatchFailCondition:
+				if stats.Counters.IDmismatch > 0 {
+					os.Exit(1)
+				}
+			}
+		}
+	}
 }
 
 func getSupportedDNSTypes() []string {
diff --git a/docs/failoncondition.md b/docs/failoncondition.md
new file mode 100644
index 0000000..d341b32
--- /dev/null
+++ b/docs/failoncondition.md
@@ -0,0 +1,29 @@
+---
+title: Fail on condition
+layout: default
+parent: Examples
+---
+
+
+
+# Fail on condition
+v3.1.0
+{: .label .label-yellow }
+*dnspyre* by default always returns a zero exit code, but this behaviour can be adjusted by using `--fail <condition>` flag, 
+which can be used to specify predefined conditions that will cause a *dnspyre* to return a non-zero exit code.
+
+Currently, the dnspyre supports these fail conditions:
+* `ioerror` = *dnspyre* exits with a non-zero status code if there is at least 1 IO error (*dnspyre* failed to send DNS request or receive DNS response)
+* `negative` = *dnspyre* exits with a non-zero status code if there is at least 1 negative DNS answer (`NXDOMAIN` or `NODATA` response)
+* `error` = *dnspyre* exits with a non-zero status code if there is at least 1 error DNS response (`SERVFAIL`, `FORMERR`, `REFUSED`, etc.)
+* `idmismatch` = *dnspyre* exits with a non-zero status code if there is at least 1 ID mismatch between DNS request and response
+
+So for example to return a non-zero exit code, when benchmark fails to send request or receive response you would specify `--fail ioerror` flag
+```
+dnspyre --server 1.2.3.4 google.com --fail ioerror
+```
+
+These fail conditions can be combined, this is achieved by repeating the `--flag` flag multiple times with different conditions.
+```
+dnspyre --server 8.8.8.8 nxdomain.cz  --fail ioerror --fail error --fail negative
+```
diff --git a/docs/index.md b/docs/index.md
index 8fc6766..a75fef0 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -82,6 +82,10 @@ Flags:
                                duration, the benchmark is canceled. This option is exclusive with --number option. The duration is
                                specified in GO duration format e.g. 10s, 15m, 1h.
       --[no-]progress          Controls whether the progress bar is shown. Enabled by default.
+      --fail=ioerror ...       Controls conditions upon which the dnspyre will exit with a non-zero exit code. Repeatable flag.
+                               Supported options are 'ioerror' (fail if there is at least 1 IO error), 'negative' (fail if there is at
+                               least 1 negative DNS answer), 'error' (fail if there is at least 1 error DNS response), 'idmismatch'
+                               (fail there is at least 1 ID mismatch between DNS request and response).
       --[no-]version           Show application version.
 
 Args: