From de71183363e2a9d49f5dcced18922bb084b066f0 Mon Sep 17 00:00:00 2001
From: Nate Sales <nate@natesales.net>
Date: Tue, 16 Jul 2024 00:56:01 -0400
Subject: [PATCH] feat: add nsid-only parameter

---
 cli/cli.go       |  1 +
 main.go          | 13 +++++++++++--
 output/pretty.go | 12 +++++++++---
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/cli/cli.go b/cli/cli.go
index 05e1b0c..40022e3 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -18,6 +18,7 @@ type Flags struct {
 	Reverse          bool          `short:"x" long:"reverse" description:"Reverse lookup"`
 	DNSSEC           bool          `short:"d" long:"dnssec" description:"Set the DO (DNSSEC OK) bit in the OPT record"`
 	NSID             bool          `short:"n" long:"nsid" description:"Set EDNS0 NSID opt"`
+	NSIDOnly         bool          `short:"N" long:"nsid-only" description:"Set EDNS0 NSID opt and query only for the NSID"`
 	ClientSubnet     string        `long:"subnet" description:"Set EDNS0 client subnet"`
 	Chaos            bool          `short:"c" long:"chaos" description:"Use CHAOS query class"`
 	Class            uint16        `short:"C" description:"Set query class (default: IN 0x01)" default:"1"`
diff --git a/main.go b/main.go
index 154b9ff..40cf2a1 100644
--- a/main.go
+++ b/main.go
@@ -368,6 +368,10 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation
 		opts.Class = dns.ClassCHAOS
 	}
 
+	if opts.NSIDOnly {
+		opts.NSID = true
+	}
+
 	// Create TLS config
 	tlsConfig := &tls.Config{
 		InsecureSkipVerify: opts.TLSInsecureSkipVerify,
@@ -499,8 +503,13 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation
 			Opts: &opts,
 		}
 
-		if opts.NSID && (opts.Format == output.FormatPretty || opts.Format == output.FormatColumn) {
-			printer.PrettyPrintNSID(entries)
+		if (opts.NSID && (opts.Format == output.FormatPretty || opts.Format == output.FormatColumn)) || opts.NSIDOnly {
+			printer.PrettyPrintNSID(entries, !opts.NSIDOnly)
+		}
+
+		// Skip printing if NSIDOnly is set
+		if opts.NSIDOnly {
+			errChan <- nil
 		}
 
 		switch opts.Format {
diff --git a/output/pretty.go b/output/pretty.go
index dd9f06a..baea21d 100644
--- a/output/pretty.go
+++ b/output/pretty.go
@@ -17,7 +17,7 @@ import (
 )
 
 // PrettyPrintNSID prints the NSID from a slice of entries
-func (p Printer) PrettyPrintNSID(entries []*Entry) {
+func (p Printer) PrettyPrintNSID(entries []*Entry, printPrefix bool) {
 	for _, entry := range entries {
 		for _, r := range entry.Replies {
 			for _, o := range r.Extra {
@@ -33,8 +33,14 @@ func (p Printer) PrettyPrintNSID(entries []*Entry) {
 							if len(entries) > 1 {
 								suffix = fmt.Sprintf(" (%s)", entry.Server)
 							}
-							util.MustWritef(p.Out, "%s %s%s\n",
-								util.Color(util.ColorWhite, "NSID:"),
+
+							var prefix string
+							if printPrefix {
+								prefix = util.Color(util.ColorWhite, "NSID:") + " "
+							}
+
+							util.MustWritef(p.Out, "%s%s%s\n",
+								prefix,
 								util.Color(util.ColorPurple, string(nsidStr)),
 								suffix,
 							)