jarvis-bot/jarvis/__init__.py

import asyncio
from datetime import datetime, timedelta
from pathlib import Path

import pymongo
from discord import DMChannel, Intents, Member, Message
from discord.ext import commands
from discord.ext.tasks import loop
from discord.utils import find, get
from discord_slash import SlashCommand
from psutil import Process

from jarvis import logo, utils
from jarvis.config import get_config
from jarvis.utils.db import DBManager

if asyncio.get_event_loop().is_closed():
    asyncio.set_event_loop(asyncio.new_event_loop())

intents = Intents.default()
intents.members = True
restart_ctx = None


jarvis = commands.Bot(command_prefix=utils.get_prefix, intents=intents)
slash = SlashCommand(jarvis, sync_commands=True, sync_on_cog_reload=True)
jarvis_self = Process()
__version__ = "0.6.0"


@jarvis.event
async def on_ready():
    global restart_ctx
    print("    Logged in as {0.user}".format(jarvis))
    print("    Connected to {} guild(s)".format(len(jarvis.guilds)))
    with jarvis_self.oneshot():
        print(f"    Current PID: {jarvis_self.pid}")
        Path(f"jarvis.{jarvis_self.pid}.pid").touch()
    if restart_ctx:
        channel = None
        if "guild" in restart_ctx:
            guild = find(lambda x: x.id == restart_ctx["guild"], jarvis.guilds)
            if guild:
                channel = find(
                    lambda x: x.id == restart_ctx["channel"], guild.channels
                )
        elif "user" in restart_ctx:
            channel = jarvis.get_user(restart_ctx["user"])
        if channel:
            await channel.send("Core systems restarted and back online.")
        restart_ctx = None


@jarvis.event
async def on_member_join(user: Member):
    guild = user.guild
    db = DBManager(get_config().mongo).mongo
    mutes = list(
        db.jarvis.mutes.find(
            {"active": True, "user": user.id, "guild": guild.id}
        )
    )
    if mutes and len(mutes) >= 1:
        mute_role = db.jarvis.settings.find_one(
            {"guild": guild.id, "setting": "mute"}
        )
        role = guild.get_role(mute_role["role"])
        await user.add_roles(
            role, reason="User is muted still muted from prior mute"
        )


@jarvis.event
async def on_message(message: Message):
    if not isinstance(message.channel, DMChannel):
        db = DBManager(get_config().mongo).mongo
        autoreact = db.jarvis.autoreact.find_one(
            {"guild": message.guild.id, "channel": message.channel.id}
        )
        if autoreact:
            for reaction in autoreact["reactions"]:
                await message.add_reaction(reaction)
        massmention = db.jarvis.settings.find_one(
            {"guild": message.guild.id, "setting": "massmention"}
        )
        if (
            massmention["value"] > 0
            and len(message.mentions) > massmention["value"]
        ):
            await message.delete()
            # TODO: Implement warnings on massmention
            await message.channel.send("Please do not mass mention.")
    await jarvis.process_commands(message)


@jarvis.event
async def on_guild_join(guild):
    general = find(lambda x: x.name == "general", guild.channels)
    if general and general.permissions_for(guild.me).send_messages:
        await general.send(
            "Allow me to introduce myself. I am J.A.R.V.I.S., a virtual "
            + "artificial intelligence, and I'm here to assist you with a "
            + "variety of tasks as best I can, "
            + "24 hours a day, seven days a week."
        )
        await asyncio.sleep(1)
        await general.send("Importing all preferences from home interface...")

        # Set some default settings
        db = DBManager(get_config().mongo).mongo
        db.jarvis.settings.insert_one(
            {"guild": guild.id, "setting": "massmention", "value": 5}
        )

        await general.send("Systems are now fully operational")


@loop(minutes=1)
async def unmute():
    db = DBManager(get_config().mongo).mongo
    mutes = list(db.jarvis.mutes.find({"active": True, "length": {"$gt": 0}}))
    mute_roles = list(
        db.jarvis.settings.find({"setting": "mute"}).sort(
            [("guild", pymongo.ASCENDING)]
        )
    )
    updates = []
    for mute in mutes:
        if mute["time"] + timedelta(minutes=mute["length"]) < datetime.now():
            mute_role = [
                x["value"] for x in mute_roles if x["guild"] == mute["guild"]
            ][0]
            guild = await jarvis.fetch_guild(mute["guild"])
            role = guild.get_role(mute_role)
            user = await guild.fetch_member(mute["user"])
            if user:
                if role in user.roles:
                    await user.remove_roles(role, reason="Unmute")
            updates.append(
                pymongo.UpdateOne(
                    {"user": user.id, "guild": guild.id, "time": mute["time"]},
                    {"$set": {"active": False}},
                )
            )
    if updates:
        db.jarvis.mutes.bulk_write(updates)


@loop(minutes=10)
async def unban():
    db = DBManager(get_config().mongo).mongo
    bans = list(db.jarvis.bans.find({"active": True, "type": "temp"}))
    updates = []
    for ban in bans:
        if ban["time"] + timedelta(
            hours=ban["length"]
        ) < datetime.now() + timedelta(minutes=10):
            guild = await jarvis.fetch_guild(ban["guild"])
            user = await jarvis.fetch_user(ban["user"])
            if user:
                guild.unban(user)
            updates.append(
                pymongo.UpdateOne(
                    {
                        "user": user.id,
                        "guild": guild.id,
                        "time": ban["time"],
                        "type": "temp",
                    },
                    {"$set": {"active": False}},
                )
            )
    if updates:
        db.jarvis.bans.bulk_write(updates)


def run(ctx=None):
    global restart_ctx
    if ctx:
        restart_ctx = ctx
    for extension in utils.get_extensions():
        jarvis.load_extension(extension)
    config = get_config()
    print(
        "    https://discord.com/api/oauth2/authorize?client_id="
        + "{}&permissions=8&scope=bot%20applications.commands".format(
            config.client_id
        )
    )
    unmute.start()
    unban.start()
    jarvis.run(config.token, bot=True, reconnect=True)
    for cog in jarvis.cogs:
        session = getattr(cog, "_session", None)
        if session:
            session.close()
    if restart_ctx:
        return restart_ctx