Fix Do{T,H,Q}

This commit is contained in:
Nexus 2024-09-18 18:09:58 +01:00
parent 820594b089
commit 2c0e61b831

View file

@ -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)