From 25fdf7b402e35712b64fb86a7cc21d84a7149342 Mon Sep 17 00:00:00 2001 From: Seaswimmer Date: Tue, 28 May 2024 18:31:38 -0400 Subject: [PATCH] feat(seautils): default to nslookup if dig is not present --- seautils/seautils.py | 94 ++++++++++++++++++++++++++------------------ 1 file changed, 55 insertions(+), 39 deletions(-) diff --git a/seautils/seautils.py b/seautils/seautils.py index 3a15c05..8ef8e1f 100644 --- a/seautils/seautils.py +++ b/seautils/seautils.py @@ -101,43 +101,59 @@ class SeaUtils(commands.Cog): 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) + if stderr: + await ctx.maybe_send_embed(message="An error was encountered!\n" + cf.box(text=stderr.decode())) else: - await ctx.send(content=cf.box(text=stdout, lang='yaml')) + 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')) + 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."))