SeaCogs/seautils/seautils.py

#   _____                         _
#  / ____|                       (_)
# | (___   ___  __ _ _____      ___ _ __ ___  _ __ ___   ___ _ __
#  \___ \ / _ \/ _` / __\ \ /\ / / | '_ ` _ \| '_ ` _ \ / _ \ '__|
#  ____) |  __/ (_| \__ \\ V  V /| | | | | | | | | | | |  __/ |
# |_____/ \___|\__,_|___/ \_/\_/ |_|_| |_| |_|_| |_| |_|\___|_|

import asyncio
import inspect
import operator
from asyncio.subprocess import Process
from functools import partial, partialmethod
from typing import Any

import yaml
from discord import Color, Embed, app_commands
from discord.utils import CachedSlotProperty, cached_property
from redbot.core import commands
from redbot.core.bot import Red
from redbot.core.dev_commands import cleanup_code
from redbot.core.utils import chat_formatting as cf
from redbot.core.utils.views import SimpleMenu


class SeaUtils(commands.Cog):
    """A collection of random utilities."""

    __author__ = ["SeaswimmerTheFsh"]
    __version__ = "1.0.0"

    def __init__(self, bot: Red):
        self.bot = bot

    def format_help_for_context(self, ctx: commands.Context) -> str:
        pre_processed = super().format_help_for_context(ctx) or ""
        n = "\n" if "\n\n" not in pre_processed else ""
        text = [
            f"{pre_processed}{n}",
            f"Cog Version: **{self.__version__}**",
            f"Author: {cf.humanize_list(self.__author__)}"
        ]
        return "\n".join(text)

    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(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(src)

    @commands.command(aliases=["source", "src", "code", "showsource"])
    @commands.is_owner()
    async def showcode(self, ctx: commands.Context, *, object: str): # pylint: disable=redefined-builtin
        """Show the code for a particular object."""
        try:
            if object.startswith("/") and (obj := ctx.bot.tree.get_command(object[1:])):
                text = self.format_src(obj)
            elif obj := ctx.bot.get_cog(object):
                text = self.format_src(type(obj))
            elif obj := ctx.bot.get_command(object):
                text = self.format_src(obj)
            temp_content = cf.pagify(
                text=cleanup_code(text),
                escape_mass_mentions=True,
                page_length = 1977
            )
            content = []
            max_i = operator.length_hint(temp_content)
            i = 1
            for page in temp_content:
                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))

    @commands.command(name='dig', aliases=['dnslookup', 'nslookup'])
    @commands.is_owner()
    async def dig(self, ctx: commands.Context, name: str, type: str | None = 'A', server: str | None = None, port: int = 53) -> None:
        """Retrieve DNS information for a domain."""
        command_opts: list[str | int] = ['dig']
        if server:
            command_opts.extend(['@', server])
        command_opts.extend([name, type])
        if port != 53:
            command_opts.extend(['-p', port])
        command_opts.extend(['+yaml'])

        try:
            process: Process = await asyncio.create_subprocess_exec(*command_opts, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
            stdout, stderr = await process.communicate()
        except (FileNotFoundError):
            await ctx.maybe_send_embed(message="The `dig` command is not installed on this system.")
            return

        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.")

                question_section = "\n".join(response_data['QUESTION_SECTION'])
                embed.add_field(name="Question Section", value=f"```{question_section}```", inline=False)

                if 'ANSWER_SECTION' in response_data:
                    answer_section = "\n".join(response_data['ANSWER_SECTION'])
                    embed.add_field(name="Answer Section", value=f"```{answer_section}```", inline=False)

                if 'AUTHORITY_SECTION' in response_data:
                    authority_section = "\n".join(response_data['AUTHORITY_SECTION'])
                    embed.add_field(name="Authority Section", value=f"```{authority_section}```", inline=False)
                await ctx.send(embed=embed)
            else:
                await ctx.send(content=cf.box(text=stdout, lang='yaml'))