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
|
import time
|
||||||
|
from email.policy import default
|
||||||
from functools import partial
|
from functools import partial
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
@ -22,6 +23,7 @@ class DNSCommand(niobot.Module):
|
||||||
dns_server: str = "internal",
|
dns_server: str = "internal",
|
||||||
):
|
):
|
||||||
"""Performs a DNS lookup and returns the result for you"""
|
"""Performs a DNS lookup and returns the result for you"""
|
||||||
|
port = 53
|
||||||
query_type = query_type.upper()
|
query_type = query_type.upper()
|
||||||
supported = ("A", "AAAA", "MX", "CNAME", "TXT", "PTR")
|
supported = ("A", "AAAA", "MX", "CNAME", "TXT", "PTR")
|
||||||
if query_type not in supported:
|
if query_type not in supported:
|
||||||
|
@ -43,7 +45,6 @@ class DNSCommand(niobot.Module):
|
||||||
elif not parsed_resolver.hostname:
|
elif not parsed_resolver.hostname:
|
||||||
return await ctx.respond("Invalid URL (no hostname)")
|
return await ctx.respond("Invalid URL (no hostname)")
|
||||||
|
|
||||||
real_resolver_url = parsed_resolver.geturl()
|
|
||||||
try:
|
try:
|
||||||
ip = ipaddress.ip_address(parsed_resolver.hostname)
|
ip = ipaddress.ip_address(parsed_resolver.hostname)
|
||||||
if ip.is_private and not self.bot.is_owner(ctx.message.sender):
|
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:
|
match parsed_resolver.scheme:
|
||||||
case "https":
|
case "https":
|
||||||
resolver_call = dns.asyncquery.https
|
resolver_call = dns.asyncquery.https
|
||||||
|
default_port = 443
|
||||||
case "quic":
|
case "quic":
|
||||||
resolver_call = dns.asyncquery.quic
|
resolver_call = dns.asyncquery.quic
|
||||||
|
default_port = 443
|
||||||
case "tls":
|
case "tls":
|
||||||
resolver_call = dns.asyncquery.tls
|
resolver_call = dns.asyncquery.tls
|
||||||
|
default_port = 853
|
||||||
case "udp":
|
case "udp":
|
||||||
resolver_call = dns.asyncquery.udp
|
resolver_call = dns.asyncquery.udp
|
||||||
|
default_port = 53
|
||||||
case "tcp":
|
case "tcp":
|
||||||
resolver_call = dns.asyncquery.tcp
|
resolver_call = dns.asyncquery.tcp
|
||||||
|
default_port = 53
|
||||||
case _:
|
case _:
|
||||||
resolver_call = None
|
resolver_call = default_port = None
|
||||||
|
|
||||||
if resolver_call is None:
|
if resolver_call is None:
|
||||||
return await ctx.respond(
|
return await ctx.respond(
|
||||||
"Unrecognised scheme %r" % parsed_resolver.scheme
|
"Unrecognised scheme %r" % parsed_resolver.scheme
|
||||||
)
|
)
|
||||||
where = real_resolver_url
|
where = parsed_resolver.hostname
|
||||||
|
port = parsed_resolver.port or default_port
|
||||||
else:
|
else:
|
||||||
try:
|
try:
|
||||||
ip = ipaddress.ip_address(dns_server)
|
ip = ipaddress.ip_address(dns_server)
|
||||||
|
@ -105,7 +112,7 @@ class DNSCommand(niobot.Module):
|
||||||
start_ts = time.perf_counter()
|
start_ts = time.perf_counter()
|
||||||
response: (
|
response: (
|
||||||
dns.message.Message | tuple[dns.message.Message, bool]
|
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()
|
end_ts = time.perf_counter()
|
||||||
if isinstance(response, dns.message.Message):
|
if isinstance(response, dns.message.Message):
|
||||||
response = (response, False)
|
response = (response, False)
|
||||||
|
|
Loading…
Reference in a new issue