radical-bot/ext/system.py

import asyncio

import lightbulb
import hikari

from lib.config import load_config

config = load_config()

plugin = lightbulb.Plugin("SystemPlugin")


async def create_subprocess(*args):
    proc = await asyncio.create_subprocess_exec(*args,
                                                stdout=asyncio.subprocess.PIPE,
                                                stderr=asyncio.subprocess.PIPE
                                                )
    stdout, _ = await proc.communicate()
    return stdout.decode("utf-8").strip()


async def get_git_latest_commit_id():
    return await create_subprocess("git", "rev-parse", "--short", "HEAD")


async def get_git_branch():
    return await create_subprocess("git", "rev-parse", "--abbrev-ref", "HEAD")


async def get_git_remote():
    remote = await create_subprocess("git", "branch", "-vv")
    return remote.split("[")[1].split("]")[0]


async def get_git_head_diff_branch(branch: str) -> str:
    # diff HEAD to remote
    return await create_subprocess("git", "diff", "head", branch)


async def get_git_index_diff_branch(branch: str) -> str:
    # diff index to remote
    return await create_subprocess("git", "diff", branch)


async def get_git_commits_ahead_behind(branch: str, remote: str) -> tuple[int, int]:
    # commits ahead/behind remote
    output = await create_subprocess("git", "rev-list", "--left-right", "--count", branch+"..."+remote)
    commits_ahead = int(output.split()[0])
    commits_behind = int(output.split()[1])
    return commits_ahead, commits_behind


async def get_git_dev(remote: str, branch: str) -> bool:
    # dev environment tests
    # difference between HEAD and remote
    if await get_git_head_diff_branch(remote) != "":
        return True
    elif config.git_branch != branch:  # branch has been changed
        return True
    # difference between index and remote
    elif await get_git_index_diff_branch(remote) != "":
        return True
    # ahead in commits
    elif (await get_git_commits_ahead_behind(branch, remote))[0] > 0:
        return True
    return False


async def get_git_update(dev_status: bool, branch, remote) -> bool:
    commits_behind, _ = await get_git_commits_ahead_behind(branch, remote)
    if not dev_status and commits_behind > 0:
        return True
    return False


async def get_git_status() -> dict:
    output = {}
    output["commit_id"] = await get_git_latest_commit_id()
    output["branch"] = await get_git_branch()
    output["remote"] = await get_git_remote()
    output["dev"] = await get_git_dev(output['remote'], output['branch'])
    output['update'] = await get_git_update(output['dev'], output['branch'], output['remote'])
    return output


@plugin.command
@lightbulb.command("ping", "pong!", ephemeral=True)
@lightbulb.implements(lightbulb.SlashCommand)
async def ping(ctx: lightbulb.Context) -> None:
    embed = hikari.Embed(title="Pong!",
                         description=f"latency: {round(ctx.bot.heartbeat_latency * 1000, ndigits=2)}ms"
                         )
    await ctx.respond(embed)


@plugin.command
@lightbulb.add_checks(lightbulb.owner_only)
@lightbulb.command("update", "update the bot!", ephemeral=True)
@lightbulb.implements(lightbulb.SlashCommand)
async def update(ctx: lightbulb.Context) -> None:
    if (await get_git_status())['dev'] is False:
        await create_subprocess("git", "pull")
    embed = hikari.Embed(title="Restarting",
                         description="Restarting to load an update!"
                         )
    await ctx.respond(embed)
    await ctx.bot.close()


@plugin.command
@lightbulb.add_checks(lightbulb.owner_only)
@lightbulb.command("branch", "get or set the working branch the bot uses.", ephemeral=True)
@lightbulb.implements(lightbulb.SlashCommandGroup)
async def branch(ctx: lightbulb.Context) -> None:
    output = await create_subprocess("git", "rev-parse", "--symbolic-full-name", "--abbrev-ref", "HEAD")
    embed = hikari.Embed(title="Current Branch",
                         description=f"Currently on branch '{output}'")
    await ctx.respond(embed=embed)


@branch.child
@lightbulb.option("name", "name of the branch", type=str, required=True)
@lightbulb.command("switch", "switch branches")
@lightbulb.implements(lightbulb.SlashCommand)
async def switch(ctx: lightbulb.Context) -> None:
    embed = hikari.Embed(title="Restarting",
                         description="Restarting to switch branches!"
                         )
    output = await create_subprocess("git", "switch", ctx.options.branch)
    if "invalid reference" in output:  # parsing output allows us more specificity.
        embed.title = "Branch does not exist."
        embed.description = "Please check your spelling."
        await ctx.respond(embed)
        return
    elif "Already on" in output:
        embed.title = f"Already on branch '{ctx.options.branch}'"
        embed.description = "Just double checking, I presume?"
        await ctx.respond(embed)
        return
    await ctx.respond(embed)
    await ctx.bot.close()


@plugin.command
@lightbulb.command("info", "get bot information such as the version, and repository link.", ephemeral=True)
@lightbulb.implements(lightbulb.SlashCommand)
async def info(ctx: lightbulb.Context) -> None:
    git_status = await get_git_status()
    embed = hikari.Embed(title="About Me!")
    embed.add_field("Git Repository", config.git_url)
    embed.add_field("Version", git_status['commit_id'], inline=True)
    embed.add_field("Branch", git_status['branch'], inline=True)
    embed.add_field("Remote", git_status['remote'], inline=True)
    embed.add_field("Needs Update?", git_status['update'])
    await ctx.respond(embed)


def load(bot: lightbulb.BotApp):
    bot.add_plugin(plugin)


def unload(bot: lightbulb.BotApp):
    bot.remove_plugin(plugin)