jarvis-bot/jarvis/cogs/botutil.py

120 lines
4.7 KiB
Python

"""JARVIS bot utility commands."""
import logging
import platform
from io import BytesIO
import psutil
from aiofile import AIOFile, LineReader
from naff import Client, Extension, PrefixedContext, prefixed_command
from naff.models.discord.embed import EmbedField
from naff.models.discord.file import File
from rich.console import Console
from jarvis.utils import build_embed
from jarvis.utils.updates import update
class BotutilCog(Extension):
"""JARVIS Bot Utility Cog."""
def __init__(self, bot: Client):
self.bot = bot
self.logger = logging.getLogger(__name__)
self.add_ext_check(self.is_owner)
async def is_owner(self, ctx: PrefixedContext) -> bool:
"""Checks if author is bot owner."""
return ctx.author.id == self.bot.owner.id
@prefixed_command(name="tail")
async def _tail(self, ctx: PrefixedContext, count: int = 10) -> None:
lines = []
async with AIOFile("jarvis.log", "r") as af:
async for line in LineReader(af):
lines.append(line)
if len(lines) == count + 1:
lines.pop(0)
log = "".join(lines)
if len(log) > 1500:
with BytesIO() as file_bytes:
file_bytes.write(log.encode("UTF8"))
file_bytes.seek(0)
log = File(file_bytes, file_name=f"tail_{count}.log")
await ctx.reply(content=f"Here's the last {count} lines of the log", file=log)
else:
await ctx.reply(content=f"```\n{log}\n```")
@prefixed_command(name="log")
async def _log(self, ctx: PrefixedContext) -> None:
async with AIOFile("jarvis.log", "r") as af:
with BytesIO() as file_bytes:
raw = await af.read_bytes()
file_bytes.write(raw)
file_bytes.seek(0)
log = File(file_bytes, file_name="jarvis.log")
await ctx.reply(content="Here's the latest log", file=log)
@prefixed_command(name="crash")
async def _crash(self, ctx: PrefixedContext) -> None:
raise Exception("As you wish")
@prefixed_command(name="sysinfo")
async def _sysinfo(self, ctx: PrefixedContext) -> None:
st_ts = int(self.bot.start_time.timestamp())
ut_ts = int(psutil.boot_time())
fields = (
EmbedField(name="Operation System", value=platform.system() or "Unknown", inline=False),
EmbedField(name="Version", value=platform.release() or "N/A", inline=False),
EmbedField(name="System Start Time", value=f"<t:{ut_ts}:F> (<t:{ut_ts}:R>)"),
EmbedField(name="Python Version", value=platform.python_version()),
EmbedField(name="Bot Start Time", value=f"<t:{st_ts}:F> (<t:{st_ts}:R>)"),
)
embed = build_embed(title="System Info", description="", fields=fields)
embed.set_image(url=self.bot.user.avatar.url)
await ctx.send(embed=embed)
@prefixed_command(name="update")
async def _update(self, ctx: PrefixedContext) -> None:
status = await update(self.bot)
if status:
console = Console()
with console.capture() as capture:
console.print(status.table)
self.logger.debug(capture.get())
self.logger.debug(len(capture.get()))
added = "\n".join(status.added)
removed = "\n".join(status.removed)
changed = "\n".join(status.changed)
fields = [
EmbedField(name="Old Commit", value=status.old_hash),
EmbedField(name="New Commit", value=status.new_hash),
]
if added:
fields.append(EmbedField(name="New Modules", value=f"```\n{added}\n```"))
if removed:
fields.append(EmbedField(name="Removed Modules", value=f"```\n{removed}\n```"))
if changed:
fields.append(EmbedField(name="Changed Modules", value=f"```\n{changed}\n```"))
embed = build_embed(
"Update Status", description="Updates have been applied", fields=fields
)
embed.set_thumbnail(url="https://dev.zevaryx.com/git.png")
self.logger.info("Updates applied")
content = f"```ansi\n{capture.get()}\n```"
if len(content) < 3000:
await ctx.reply(content, embed=embed)
else:
await ctx.reply(f"Total Changes: {status.lines['total_lines']}", embed=embed)
else:
embed = build_embed(title="Update Status", description="No changes applied", fields=[])
embed.set_thumbnail(url="https://dev.zevaryx.com/git.png")
await ctx.reply(embed=embed)
def setup(bot: Client) -> None:
"""Add BotutilCog to JARVIS"""
BotutilCog(bot)