From 2c0e61b8319e1a5fe04d927fee5fea0cef7f30bc Mon Sep 17 00:00:00 2001
From: nexy7574 <git@nexy7574.co.uk>
Date: Wed, 18 Sep 2024 18:09:58 +0100
Subject: [PATCH] Fix Do{T,H,Q}

---
 app/modules/drill.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/app/modules/drill.py b/app/modules/drill.py
index 38d0c75..0b7bc38 100644
--- a/app/modules/drill.py
+++ b/app/modules/drill.py
@@ -1,4 +1,5 @@
 import time
+from email.policy import default
 from functools import partial
 
 import httpx
@@ -22,6 +23,7 @@ class DNSCommand(niobot.Module):
         dns_server: str = "internal",
     ):
         """Performs a DNS lookup and returns the result for you"""
+        port = 53
         query_type = query_type.upper()
         supported = ("A", "AAAA", "MX", "CNAME", "TXT", "PTR")
         if query_type not in supported:
@@ -43,7 +45,6 @@ class DNSCommand(niobot.Module):
             elif not parsed_resolver.hostname:
                 return await ctx.respond("Invalid URL (no hostname)")
 
-            real_resolver_url = parsed_resolver.geturl()
             try:
                 ip = ipaddress.ip_address(parsed_resolver.hostname)
                 if ip.is_private and not self.bot.is_owner(ctx.message.sender):
@@ -56,22 +57,28 @@ class DNSCommand(niobot.Module):
             match parsed_resolver.scheme:
                 case "https":
                     resolver_call = dns.asyncquery.https
+                    default_port = 443
                 case "quic":
                     resolver_call = dns.asyncquery.quic
+                    default_port = 443
                 case "tls":
                     resolver_call = dns.asyncquery.tls
+                    default_port = 853
                 case "udp":
                     resolver_call = dns.asyncquery.udp
+                    default_port = 53
                 case "tcp":
                     resolver_call = dns.asyncquery.tcp
+                    default_port = 53
                 case _:
-                    resolver_call = None
+                    resolver_call = default_port = None
 
             if resolver_call is None:
                 return await ctx.respond(
                     "Unrecognised scheme %r" % parsed_resolver.scheme
                 )
-            where = real_resolver_url
+            where = parsed_resolver.hostname
+            port = parsed_resolver.port or default_port
         else:
             try:
                 ip = ipaddress.ip_address(dns_server)
@@ -105,7 +112,7 @@ class DNSCommand(niobot.Module):
             start_ts = time.perf_counter()
             response: (
                 dns.message.Message | tuple[dns.message.Message, bool]
-            ) = await resolver_call(q=query, where=where, timeout=60)
+            ) = await resolver_call(q=query, where=where, port=port, timeout=60)
             end_ts = time.perf_counter()
             if isinstance(response, dns.message.Message):
                 response = (response, False)