nex/nonsensebot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Do{T,H,Q}

Browse source
This commit is contained in:
Nexus 2024-09-18 18:09:58 +01:00
parent 820594b089
commit 2c0e61b831
1 changed files with 11 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
app/modules/drill.py
View file

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