jarvis-bot/jarvis/cogs/botutil.py
2022-05-02 01:57:35 -06:00

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 dis_snek import PrefixedContext, Scale, Snake, prefixed_command
from dis_snek.client.errors import HTTPException
from dis_snek.models.discord.embed import EmbedField
from dis_snek.models.discord.file import File
from rich.console import Console
from jarvis.utils import build_embed
from jarvis.utils.updates import update
class BotutilCog(Scale):
"""JARVIS Bot Utility Cog."""
def __init__(self, bot: Snake):
self.bot = bot
self.logger = logging.getLogger(__name__)
self.add_scale_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")
try:
await ctx.reply(f"```ansi\n{capture.get()}\n```", embed=embed)
except HTTPException:
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: Snake) -> None:
"""Add BotutilCog to JARVIS"""
BotutilCog(bot)