From 01b03df8934ea508c285616e92a562803a70ead9 Mon Sep 17 00:00:00 2001
From: Chris <chrisbreuer93@gmail.com>
Date: Sat, 18 Jan 2025 23:55:08 +0100
Subject: [PATCH] feat: allow for querying multiple records

---
 bin/cli.ts    | 47 +++++++++++++++++++++++++++++++----------------
 src/output.ts |  2 +-
 2 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/bin/cli.ts b/bin/cli.ts
index b388bdd..4b0df7d 100644
--- a/bin/cli.ts
+++ b/bin/cli.ts
@@ -37,7 +37,7 @@ export interface CliOptions extends DnsOptions {
 const cli: CAC = cac('dnsx')
 
 cli
-  .command('[...domains]', 'Perform DNS lookup for specified domains')
+  .command('[...args]', 'Perform DNS lookup for specified domains')
   .option('-q, --query <HOST>', 'Host name or domain name to query')
   .option('-t, --type <TYPE>', 'Type of the DNS record being queried (A, MX, NS...)')
   .option('-n, --nameserver <ADDR>', 'Address of the nameserver to send packets to')
@@ -57,36 +57,51 @@ cli
   .option('--verbose', 'Print additional debugging information', { default: false })
   .example('dnsx example.com')
   .example('dnsx example.com MX')
-  .example('dnsx example.com MX @1.1.1.1')
+  .example('dnsx example.com A AAAA NS MX')
   .example('dnsx example.com -t MX -n 1.1.1.1 -T')
-  .action(async (domains: string[], options: CliOptions) => {
+  .action(async (args: string[], options: CliOptions) => {
     try {
-      // Handle domains from both arguments and --query option
-      const allDomains = [
-        ...domains,
-      ]
-
-      // If we have a non-option argument after domain that looks like a record type, use it
-      if (domains.length > 1 && /^[A-Z]+$/i.test(domains[1])) {
-        options.type = domains[1]
-        allDomains.pop() // Remove the type from domains array
+      // Separate domains and record types from arguments
+      const validRecordTypes = new Set(['A', 'AAAA', 'NS', 'MX', 'TXT', 'SRV', 'PTR', 'CNAME', 'SOA', 'CAA'])
+      const domains: string[] = []
+      const types: string[] = []
+
+      // First argument is always a domain
+      if (args.length > 0) {
+        domains.push(args[0])
+      }
+
+      // Remaining arguments could be record types
+      for (let i = 1; i < args.length; i++) {
+        const arg = args[i].toUpperCase()
+        if (validRecordTypes.has(arg)) {
+          types.push(arg)
+        } else {
+          domains.push(arg)
+        }
       }
 
       // Add any domains from --query option
       if (options.query) {
-        allDomains.push(...(Array.isArray(options.query) ? options.query : [options.query]))
+        domains.push(...(Array.isArray(options.query) ? options.query : [options.query]))
+      }
+
+      // Add any types from -t/--type option
+      if (options.type) {
+        types.push(...(Array.isArray(options.type) ? options.type : [options.type]))
       }
 
       // Check if we have any domains to query
-      if (allDomains.length === 0) {
+      if (domains.length === 0) {
+        console.error(colors.error('Error: No domains specified'))
         process.exit(1)
       }
 
       // Create client with parsed options
       const client = new DnsClient({
-        domains: allDomains,
+        domains,
         nameserver: options.nameserver,
-        type: options.type,
+        type: types.length > 0 ? types : undefined, // Only set if we have types
         class: options.class,
         udp: options.udp,
         tcp: options.tcp,
diff --git a/src/output.ts b/src/output.ts
index 5b5751c..b7afd2c 100644
--- a/src/output.ts
+++ b/src/output.ts
@@ -123,7 +123,7 @@ function formatText(responses: DnsResponse[], options: OutputOptions): string {
 
     // Full mode - all sections
     if (response.answers.length > 0) {
-      output.push('\nAnswers:')
+      // output.push('\nAnswers:')
       for (const answer of response.answers) {
         output.push(formatAnswer(answer, options))
       }