Fix Do{T,H,Q}
This commit is contained in:
parent
820594b089
commit
2c0e61b831
1 changed files with 11 additions and 4 deletions
|
@ -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)
|
||||
|
|
Loading…
Reference in a new issue