SeaCogs/seautils/seautils.py

255 lines
14 KiB
Python
Raw Permalink Normal View History

2024-05-13 19:26:13 -04:00
# _____ _
# / ____| (_)
# | (___ ___ __ _ _____ ___ _ __ ___ _ __ ___ ___ _ __
# \___ \ / _ \/ _` / __\ \ /\ / / | '_ ` _ \| '_ ` _ \ / _ \ '__|
# ____) | __/ (_| \__ \\ V V /| | | | | | | | | | | | __/ |
# |_____/ \___|\__,_|___/ \_/\_/ |_|_| |_| |_|_| |_| |_|\___|_|
2024-05-28 16:22:22 -04:00
import asyncio
2024-05-13 19:26:13 -04:00
import inspect
2024-05-13 21:12:51 -04:00
import operator
import re
from asyncio.subprocess import Process
from functools import partial, partialmethod
from typing import Any
2024-05-13 19:26:13 -04:00
2024-05-28 20:20:21 -04:00
import aiohttp
import yaml
2024-05-28 20:20:21 -04:00
from bs4 import BeautifulSoup
from discord import Color, Embed, app_commands
from discord.utils import CachedSlotProperty, cached_property
from markdownify import MarkdownConverter
2024-05-13 19:26:13 -04:00
from redbot.core import commands
from redbot.core.bot import Red
from redbot.core.dev_commands import cleanup_code
2024-05-13 19:26:13 -04:00
from redbot.core.utils import chat_formatting as cf
from redbot.core.utils.views import SimpleMenu
2024-05-13 19:26:13 -04:00
def md(soup: BeautifulSoup, **options) -> Any | str:
2024-06-01 15:03:23 -04:00
return MarkdownConverter(**options).convert_soup(soup=soup)
def format_rfc_text(text: str, number: int) -> str:
one: str = re.sub(r"\(\.\/rfc(\d+)", r"(https://www.rfc-editor.org/rfc/rfc\1.html", text)
two: str = re.sub(r"\((#(?:section|page)-\d+(?:.\d+)?)\)", f"(https://www.rfc-editor.org/rfc/rfc{number}.html\1)", one)
three: str = re.sub(r"\n{3,}", "\n\n", two)
return three
2024-05-13 19:26:13 -04:00
class SeaUtils(commands.Cog):
"""A collection of random utilities."""
__author__ = ["[cswimr](https://www.coastalcommits.com/cswimr)"]
__git__ = "https://www.coastalcommits.com/cswimr/SeaCogs"
__version__ = "1.0.1"
__documentation__ = "https://seacogs.coastalcommits.com/seautils/"
2024-05-13 19:26:13 -04:00
def __init__(self, bot: Red) -> None:
2024-05-13 19:26:13 -04:00
self.bot = bot
def format_help_for_context(self, ctx: commands.Context) -> str:
pre_processed = super().format_help_for_context(ctx) or ""
2024-05-13 19:26:13 -04:00
n = "\n" if "\n\n" not in pre_processed else ""
text = [
f"{pre_processed}{n}",
f"{cf.bold('Cog Version:')} [{self.__version__}]({self.__git__})",
f"{cf.bold('Author:')} {cf.humanize_list(self.__author__)}",
f"{cf.bold('Documentation:')} {self.__documentation__}",
2024-05-13 19:26:13 -04:00
]
return "\n".join(text)
2024-05-19 00:18:48 -04:00
def format_src(self, obj: Any) -> str:
"""A large portion of this code is repurposed from Zephyrkul's RTFS cog.
https://github.com/Zephyrkul/FluffyCogs/blob/master/rtfs/rtfs.py"""
obj = inspect.unwrap(func=obj)
src: Any = getattr(obj, "__func__", obj)
if isinstance(obj, (commands.Command, app_commands.Command)):
src = obj.callback
elif isinstance(obj, (partial, partialmethod)):
src = obj.func
elif isinstance(obj, property):
src = obj.fget
elif isinstance(obj, (cached_property, CachedSlotProperty)):
src = obj.function
return inspect.getsource(object=src)
@commands.command(aliases=["source", "src", "code", "showsource"])
2024-05-13 19:26:13 -04:00
@commands.is_owner()
async def showcode(self, ctx: commands.Context, *, object: str) -> None: # pylint: disable=redefined-builtin
"""Show the code for a particular object."""
2024-05-13 19:26:13 -04:00
try:
if object.startswith("/") and (obj := ctx.bot.tree.get_command(object[1:])):
2024-05-19 00:18:48 -04:00
text = self.format_src(obj)
elif obj := ctx.bot.get_cog(object):
2024-05-19 00:18:48 -04:00
text = self.format_src(type(obj))
2024-05-17 00:20:00 -04:00
elif obj := ctx.bot.get_command(object):
2024-05-19 00:18:48 -04:00
text = self.format_src(obj)
2024-06-04 12:16:09 -04:00
else:
raise AttributeError
2024-05-13 21:11:15 -04:00
temp_content = cf.pagify(
text=cleanup_code(text),
escape_mass_mentions=True,
page_length = 1977
2024-05-13 21:11:15 -04:00
)
content = []
2024-05-13 21:12:51 -04:00
max_i = operator.length_hint(temp_content)
i = 1
for page in temp_content:
2024-05-13 21:11:15 -04:00
content.append(f"**Page {i}/{max_i}**\n{cf.box(page, lang='py')}")
i += 1
await SimpleMenu(pages=content, disable_after_timeout=True, timeout=180).start(ctx)
except (OSError, AttributeError, UnboundLocalError):
if ctx.embed_requested():
embed = Embed(title="Object not found!", color=await ctx.embed_color())
await ctx.send(embed=embed, reference=ctx.message.to_reference(fail_if_not_exists=False))
else:
await ctx.send(content="Object not found!", reference=ctx.message.to_reference(fail_if_not_exists=False))
2024-05-28 16:22:22 -04:00
@commands.command(name='dig', aliases=['dnslookup', 'nslookup'])
@commands.is_owner()
2024-05-30 11:02:22 -04:00
async def dig(self, ctx: commands.Context, name: str, record_type: str | None = None, server: str | None = None, port: int = 53) -> None:
"""Retrieve DNS information for a domain.
Uses `dig` to perform a DNS query. Will fall back to `nslookup` if `dig` is not installed on the system.
`nslookup` does not provide as much information as `dig`, so only the `name` parameter will be used if `nslookup` is used.
Will return the A, AAAA, and CNAME records for a domain by default. You can specify a different record type with the `type` parameter."""
2024-05-28 18:40:43 -04:00
command_opts: list[str | int] = ['dig']
2024-05-30 11:02:22 -04:00
query_types: list[str] = [record_type] if record_type else ['A', 'AAAA', 'CNAME']
2024-05-28 16:22:22 -04:00
if server:
command_opts.extend(['@', server])
for query_type in query_types:
command_opts.extend([name, query_type])
command_opts.extend(['-p', str(port), '+yaml'])
2024-05-28 16:22:22 -04:00
try:
process: Process = await asyncio.create_subprocess_exec(*command_opts, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
stdout, stderr = await process.communicate()
if stderr:
await ctx.maybe_send_embed(message="An error was encountered!\n" + cf.box(text=stderr.decode()))
else:
data = yaml.safe_load(stdout.decode())
message_data: dict = data[0]['message']
response_data: dict = message_data['response_message_data']
if ctx.embed_requested():
embed = Embed(
title="DNS Query Result",
color=await ctx.embed_color(),
timestamp=message_data['response_time']
)
embed.add_field(name="Response Address", value=message_data['response_address'], inline=True)
embed.add_field(name="Response Port", value=message_data['response_port'], inline=True)
embed.add_field(name="Query Address", value=message_data['query_address'], inline=True)
embed.add_field(name="Query Port", value=message_data['query_port'], inline=True)
embed.add_field(name="Status", value=response_data['status'], inline=True)
embed.add_field(name="Flags", value=response_data['flags'], inline=True)
if response_data.get('status') != 'NOERROR':
embed.colour = Color.red()
embed.description = cf.error("Dig query did not return `NOERROR` status.")
questions = []
answers = []
authorities = []
for m in data:
response = m['message']['response_message_data']
if 'QUESTION_SECTION' in response:
for question in response['QUESTION_SECTION']:
if question not in questions:
questions.append(question)
if 'ANSWER_SECTION' in response:
for answer in response['ANSWER_SECTION']:
if answer not in answers:
answers.append(answer)
if 'AUTHORITY_SECTION' in response:
for authority in response['AUTHORITY_SECTION']:
if authority not in authorities:
authorities.append(authority)
if questions:
question_section = "\n".join(questions)
embed.add_field(name="Question Section", value=f"{cf.box(text=question_section, lang='prolog')}", inline=False)
if answers:
answer_section = "\n".join(answers)
if len(answer_section) > 1024:
embed.description = cf.warning("Answer section is too long to fit within embed field, falling back to description.") + cf.box(answer_section)
else:
embed.add_field(name="Answer Section", value=f"{cf.box(text=answer_section, lang='prolog')}", inline=False)
if authorities:
authority_section = "\n".join(authorities)
embed.add_field(name="Authority Section", value=f"{cf.box(text=authority_section, lang='prolog')}", inline=False)
await ctx.send(embed=embed)
else:
await ctx.send(content=cf.box(text=stdout, lang='yaml'))
except (FileNotFoundError):
try:
ns_process = await asyncio.create_subprocess_exec('nslookup', name, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
ns_stdout, ns_stderr = await ns_process.communicate()
if ns_stderr:
await ctx.maybe_send_embed(message="An error was encountered!\n" + cf.box(text=ns_stderr.decode()))
else:
warning = cf.warning("`dig` is not installed! Defaulting to `nslookup`.\nThis command provides more information when `dig` is installed on the system.\n")
if await ctx.embed_requested():
embed = Embed(
title="DNS Query Result",
color=await ctx.embed_color(),
timestamp=ctx.message.created_at
)
embed.description = warning + cf.box(text=ns_stdout.decode())
await ctx.send(embed=embed)
else:
await ctx.send(content = warning + cf.box(text=ns_stdout.decode()))
except (FileNotFoundError):
await ctx.maybe_send_embed(message=cf.error("Neither `dig` nor `nslookup` are installed on the system. Unable to resolve DNS query."))
2024-05-28 20:20:21 -04:00
@commands.command()
async def rfc(self, ctx: commands.Context, number: int) -> None:
"""Retrieve the text of an RFC document.
This command uses the [RFC Editor website](https://www.rfc-editor.org/) to fetch the text of an RFC document.
A [Request for Comments (RFC)](https://en.wikipedia.org/wiki/Request_for_Comments) is a publication in a series from the principal technical development and standards-setting bodies for the [Internet](https://en.wikipedia.org/wiki/Internet), most prominently the [Internet Engineering Task Force](https://en.wikipedia.org/wiki/Internet_Engineering_Task_Force). An RFC is authored by individuals or groups of engineers and [computer scientists](https://en.wikipedia.org/wiki/Computer_scientist) in the form of a [memorandum](https://en.wikipedia.org/wiki/Memorandum) describing methods, behaviors, research, or innovations applicable to the working of the Internet and Internet-connected systems. It is submitted either for [peer review](https://en.wikipedia.org/wiki/Peer_review) or to convey new concepts, information, or, occasionally, engineering humor.""" # noqa: E501
2024-05-28 20:20:21 -04:00
url = f"https://www.rfc-editor.org/rfc/rfc{number}.html"
datatracker_url = f"https://datatracker.ietf.org/doc/rfc{number}"
2024-05-28 20:20:21 -04:00
async with aiohttp.ClientSession() as session:
async with session.get(url=url) as response:
if response.status == 200:
html = await response.text()
soup = BeautifulSoup(html, 'html.parser')
pre_tags = soup.find_all('pre')
2024-05-28 21:54:48 -04:00
content: list[Embed | str] = []
for pre_tag in pre_tags:
text = format_rfc_text(md(pre_tag), number)
if len(text) > 4096:
2024-05-28 21:54:48 -04:00
pagified_text = cf.pagify(text, delims=["\n\n"], page_length=4096)
for page in pagified_text:
if await ctx.embed_requested():
embed = Embed(
title=f"RFC Document {number}",
url=datatracker_url,
2024-05-28 21:54:48 -04:00
description=page,
color=await ctx.embed_color()
)
content.append(embed)
else:
content.append(page)
else:
if await ctx.embed_requested():
embed = Embed(
title=f"RFC Document {number}",
url=datatracker_url,
description=text,
color=await ctx.embed_color()
)
content.append(embed)
else:
content.append(text)
2024-05-28 21:54:48 -04:00
if await ctx.embed_requested():
for embed in content:
embed.set_footer(text=f"Page {content.index(embed) + 1}/{len(content)}")
2024-05-28 20:20:21 -04:00
await SimpleMenu(pages=content, disable_after_timeout=True, timeout=300).start(ctx)
else:
await ctx.maybe_send_embed(message=cf.error(f"An error occurred while fetching RFC {number}. Status code: {response.status}."))