From 0038eecfbc10026cafc83f448c9d2c4053a996f9 Mon Sep 17 00:00:00 2001
From: jkhsjdhjs <jkhsjdhjs@totally.rip>
Date: Mon, 17 Jul 2023 01:27:26 +0000
Subject: [PATCH] nslookup: support lookup of multiple record types at once

---
 bert.py     |  9 ++++-----
 nslookup.py | 25 +++++++++++++------------
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/bert.py b/bert.py
index a6c8214..9b7a6f3 100755
--- a/bert.py
+++ b/bert.py
@@ -460,13 +460,12 @@ def command_loop():
                   except:
                      say(channel, r+"You must specify a domain!")
                      continue
-                  try:
-                     nstype = msg_parts[2]
-                  except:
-                     nstype = "A"
+                  nstype = msg_parts[2:]
+                  if not nstype:
+                      nstype = ["AAAA", "A"]
                   print("got nslookup command with domain: "+str(nsdomain)+" and type: "+str(nstype))
                   nsresult = nslookup(nsdomain, nstype)
-                  nslines = nsresult.split("\n")
+                  nslines = set(nsresult.split("\n"))
                   for nsline in nslines:
                      say(channel, y+" "+str(nsline))
                except Exception as err:
diff --git a/nslookup.py b/nslookup.py
index e0ab12c..767a9ee 100644
--- a/nslookup.py
+++ b/nslookup.py
@@ -2,20 +2,21 @@
 
 import dns.resolver
 
-def nslookup(domain, typ="AAAA"):
+def nslookup(domain, types=["AAAA", "A"]):
    result = ""
    rcount = 0
    try:
-      answer = dns.resolver.resolve(domain, typ)
-      for rdata in answer:
-         if rcount > 0:
-            result += "\n"
-         if hasattr(rdata, 'exchange'):
-            result += str(rdata.preference)+" "+str(rdata.exchange)
-            rcount += 1
-         else:
-            result += str(rdata)
-            rcount += 1
+      for typ in types:
+         answer = dns.resolver.resolve(domain, typ)
+         for rdata in answer:
+            if rcount > 0:
+               result += "\n"
+            if hasattr(rdata, 'exchange'):
+               result += str(rdata.preference)+" "+str(rdata.exchange)
+               rcount += 1
+            else:
+               result += str(rdata)
+               rcount += 1
    except dns.resolver.NXDOMAIN as err:
       result = str(err)
    except Exception as err:
@@ -25,7 +26,7 @@ def nslookup(domain, typ="AAAA"):
    return result
 
 def main():
-   print(nslookup("elmo.space", "AAAA"))
+   print(nslookup("elmo.space"))
 
 if __name__ == "__main__":
    main()