From 028cae9e99e2d60ccaa26d87bd1ad19f1f7fdf4b Mon Sep 17 00:00:00 2001 From: Seaswimmer Date: Tue, 28 May 2024 17:53:58 -0400 Subject: [PATCH] feat(seautils): improve error code handling --- seautils/seautils.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/seautils/seautils.py b/seautils/seautils.py index 198a748..c884d08 100644 --- a/seautils/seautils.py +++ b/seautils/seautils.py @@ -12,7 +12,7 @@ from functools import partial, partialmethod from typing import Any import yaml -from discord import Embed, app_commands +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 @@ -103,8 +103,8 @@ class SeaUtils(commands.Cog): await ctx.maybe_send_embed(content= "An error was encountered!\n" + cf.box(text=stderr.decode())) else: data = yaml.safe_load(stdout.decode()) - message_data = data[0]['message'] - response_data = message_data['response_message_data'] + message_data: dict = data[0]['message'] + response_data: dict = message_data['response_message_data'] if ctx.embed_requested(): embed = Embed( title="DNS Query Result", @@ -118,12 +118,23 @@ class SeaUtils(commands.Cog): embed.add_field(name="Status", value=response_data['status'], inline=True) embed.add_field(name="Flags", value=response_data['flags'], inline=True) + match response_data.get('status'): + case 'FORMERR', 'SERVFAIL', 'NXDOMAIN', 'NOTIMP', 'REFUSED': + embed.colour = Color.red() + embed.description = "Dig query did not return `NOERROR` status." + case _: + embed.colour = await ctx.embed_color() + 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'))