diff --git a/config.example.yaml b/config.example.yaml index b320cb0..cfa17f6 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -3,12 +3,22 @@ client_id: 123456789012345678 logo: alligator2 mongo: - username: user - password: pass - host: localhost - port: 27017 + connect: + username: user + password: pass + host: localhost + port: 27017 + database: database urls: url_name: url url_name2: url2 max_messages: 1000 gitlab_token: null + cogs: + - list + - of + - enabled + - cogs + - all + - if + - empty diff --git a/jarvis/__init__.py b/jarvis/__init__.py index a0b6b17..664f15d 100644 --- a/jarvis/__init__.py +++ b/jarvis/__init__.py @@ -4,13 +4,13 @@ from pathlib import Path from discord import Intents from discord.ext import commands from discord.utils import find +from discord_slash import SlashCommand +from psutil import Process + from jarvis import logo, tasks, utils from jarvis.config import get_config from jarvis.db import DBManager from jarvis.events import guild, member, message -from psutil import Process - -from discord_slash import SlashCommand if asyncio.get_event_loop().is_closed(): asyncio.set_event_loop(asyncio.new_event_loop()) @@ -25,11 +25,11 @@ jarvis = commands.Bot( ) slash = SlashCommand(jarvis, sync_commands=True, sync_on_cog_reload=True) jarvis_self = Process() -__version__ = "1.8.1" +__version__ = "1.9.0" - -db = DBManager(get_config().mongo).mongo -jarvis_db = db.jarvis +jconfig = get_config() +db = DBManager(jconfig.mongo["connect"]).mongo +jarvis_db = db[jconfig.mongo["database"]] @jarvis.event diff --git a/jarvis/cogs/admin/lockdown.py b/jarvis/cogs/admin/lockdown.py index 38c09ce..20a7f43 100644 --- a/jarvis/cogs/admin/lockdown.py +++ b/jarvis/cogs/admin/lockdown.py @@ -15,7 +15,10 @@ from jarvis.utils.permissions import admin_or_permissions class LockdownCog(CacheCog): def __init__(self, bot: commands.Bot): super().__init__(bot) - self.db = DBManager(get_config().mongo).mongo.jarvis + config = get_config() + self.db = DBManager(config.mongo["connect"]).mongo[ + config.mongo["database"] + ] @cog_ext.cog_subcommand( base="lockdown", diff --git a/jarvis/cogs/ctc2.py b/jarvis/cogs/ctc2.py index 7fdf00c..30dc32e 100644 --- a/jarvis/cogs/ctc2.py +++ b/jarvis/cogs/ctc2.py @@ -30,7 +30,7 @@ class CTCCog(CacheCog): def __init__(self, bot): super().__init__(bot) mconf = get_config().mongo - self.db = DBManager(mconf).mongo + self.db = DBManager(mconf["connect"]).mongo self._session = aiohttp.ClientSession() self.url = "https://completethecodetwo.cards/pw" diff --git a/jarvis/cogs/error.py b/jarvis/cogs/error.py index b52a2ba..6009ce8 100644 --- a/jarvis/cogs/error.py +++ b/jarvis/cogs/error.py @@ -28,7 +28,8 @@ class ErrorHandlerCog(commands.Cog): return elif isinstance(error, commands.errors.CommandOnCooldown): await ctx.send( - "Command on cooldown. Please wait before trying again", + "Command on cooldown. " + + f"Please wait {error.retry_after:0.2f}s before trying again", hidden=True, ) else: diff --git a/jarvis/cogs/jokes.py b/jarvis/cogs/jokes.py index c5db5ed..240bcad 100644 --- a/jarvis/cogs/jokes.py +++ b/jarvis/cogs/jokes.py @@ -8,6 +8,7 @@ from discord.ext import commands from discord_slash import cog_ext import jarvis +from jarvis.config import get_config from jarvis.db import DBManager from jarvis.utils import build_embed from jarvis.utils.field import Field @@ -22,8 +23,10 @@ class JokeCog(commands.Cog): def __init__(self, bot): self.bot = bot - config = jarvis.config.get_config() - self.db = DBManager(config.mongo).mongo.jarvis + config = get_config() + self.db = DBManager(config.mongo["connect"]).mongo[ + config.mongo["database"] + ] # TODO: Make this a command group with subcommands async def _joke(self, ctx, id: str = None): diff --git a/jarvis/cogs/modlog/command.py b/jarvis/cogs/modlog/command.py index ff1b7dd..34c07c8 100644 --- a/jarvis/cogs/modlog/command.py +++ b/jarvis/cogs/modlog/command.py @@ -13,7 +13,7 @@ class ModlogCommandCog(commands.Cog): @commands.Cog.listener() async def on_slash_command(self, ctx: SlashContext): - if not isinstance(ctx.channel, DMChannel): + if not isinstance(ctx.channel, DMChannel) and ctx.name not in ["pw"]: modlog = Setting.get(guild=ctx.guild.id, setting="modlog") if modlog: channel = ctx.guild.get_channel(modlog.value) diff --git a/jarvis/cogs/modlog/utils.py b/jarvis/cogs/modlog/utils.py index 4830c78..b0c713b 100644 --- a/jarvis/cogs/modlog/utils.py +++ b/jarvis/cogs/modlog/utils.py @@ -2,12 +2,12 @@ from datetime import datetime, timedelta import discord from discord.utils import find + from jarvis.utils import build_embed from jarvis.utils.field import Field def modlog_embed( - self, member: discord.Member, admin: discord.Member, log: discord.AuditLogEntry, diff --git a/jarvis/cogs/owner.py b/jarvis/cogs/owner.py index 7ea6cb8..c828339 100644 --- a/jarvis/cogs/owner.py +++ b/jarvis/cogs/owner.py @@ -35,7 +35,11 @@ class OwnerCog(commands.Cog): or ctx.message.author.id in self.admins ): try: + if "jarvis.cogs." not in cog: + cog = "jarvis.cogs." + cog.split(".")[-1] self.bot.load_extension(cog) + except commands.errors.ExtensionAlreadyLoaded: + await ctx.send(f"Cog `{cog}` already loaded") except Exception as e: await ctx.send( f"Failed to load new cog `{cog}`: {type(e).name} - {e}" @@ -48,7 +52,7 @@ class OwnerCog(commands.Cog): @commands.command(name="unload", hidden=True) @user_is_bot_admin() async def _unload_cog(self, ctx, *, cog: str): - if cog == "jarvis.cogs.owner": + if cog in ["jarvis.cogs.owner", "owner"]: await ctx.send("Cannot unload `owner` cog") return info = await self.bot.application_info() @@ -57,7 +61,11 @@ class OwnerCog(commands.Cog): or ctx.message.author.id in self.admins ): try: + if "jarvis.cogs." not in cog: + cog = "jarvis.cogs." + cog.split(".")[-1] self.bot.unload_extension(cog) + except commands.errors.ExtensionNotLoaded: + await ctx.send(f"Cog `{cog}` not loaded") except Exception as e: await ctx.send( f"Failed to unload cog `{cog}` {type(e).__name__} - {e}" @@ -70,7 +78,7 @@ class OwnerCog(commands.Cog): @commands.command(name="reload", hidden=True) @user_is_bot_admin() async def _cog_reload(self, ctx, *, cog: str): - if cog == "jarvis.cogs.owner": + if cog in ["jarvis.cogs.owner", "owner"]: await ctx.send("Cannot reload `owner` cog") return info = await self.bot.application_info() @@ -79,8 +87,13 @@ class OwnerCog(commands.Cog): or ctx.message.author.id in self.admins ): try: + if "jarvis.cogs." not in cog: + cog = "jarvis.cogs." + cog.split(".")[-1] + try: + self.bot.load_extension(cog) + except commands.errors.ExtensionNotLoaded: + pass self.bot.unload_extension(cog) - self.bot.load_extension(cog) except Exception as e: await ctx.send( f"Failed to reload cog `{cog}` {type(e).__name__} - {e}" diff --git a/jarvis/cogs/util.py b/jarvis/cogs/util.py index 6e5123d..0791684 100644 --- a/jarvis/cogs/util.py +++ b/jarvis/cogs/util.py @@ -1,14 +1,17 @@ import re +import secrets +import string from io import BytesIO from discord import File, Guild, Role, User from discord.ext import commands from discord_slash import SlashContext, cog_ext -from discord_slash.utils.manage_commands import create_option +from discord_slash.utils.manage_commands import create_choice, create_option from PIL import Image, ImageDraw import jarvis -from jarvis import config, jarvis_self, logo +from jarvis import jarvis_self, logo +from jarvis.config import get_config from jarvis.data.robotcamo import emotes, names from jarvis.utils import build_embed, convert_bytesize, get_repo_hash from jarvis.utils.field import Field @@ -23,7 +26,7 @@ class UtilCog(commands.Cog): def __init__(self, bot): self.bot = bot - self.config = config.get_config() + self.config = get_config() @cog_ext.cog_slash( name="status", @@ -267,6 +270,57 @@ class UtilCog(commands.Cog): await ctx.send(embed=embed) + @cog_ext.cog_subcommand( + base="pw", + name="gen", + base_desc="Password utilites", + description="Generate a secure password", + guild_ids=[862402786116763668], + options=[ + create_option( + name="length", + description="Password length (default 32)", + option_type=4, + required=False, + ), + create_option( + name="chars", + description="Characters to include (default last option)", + option_type=4, + required=False, + choices=[ + create_choice(name="A-Za-z", value=0), + create_choice(name="A-Fa-f0-9", value=1), + create_choice(name="A-Za-z0-9", value=2), + create_choice(name="A-Za-z0-9!@#$%^&*", value=3), + ], + ), + ], + ) + @commands.cooldown(1, 15, type=commands.BucketType.user) + async def _pw_gen( + self, ctx: SlashContext, length: int = 32, chars: int = 3 + ): + if length > 256: + await ctx.send( + "Please limit password to 256 characters", hidden=True + ) + return + choices = [ + string.ascii_letters, + string.hexdigits, + string.ascii_letters + string.digits, + string.ascii_letters + string.digits + "!@#$%^&*", + ] + + pw = "".join(secrets.choice(choices[chars]) for i in range(length)) + await ctx.send( + f"Generated password:\n`{pw}`\n\n" + + '**WARNING: Once you press "Dismiss Message", ' + + "*the password is lost forever***", + hidden=True, + ) + def setup(bot): bot.add_cog(UtilCog(bot)) diff --git a/jarvis/config.py b/jarvis/config.py index 3f3883d..89e426c 100644 --- a/jarvis/config.py +++ b/jarvis/config.py @@ -24,6 +24,7 @@ class Config(object): logo: str, mongo: dict, urls: dict, + cogs: list = None, gitlab_token: str = None, max_messages: int = 1000, ): @@ -32,9 +33,10 @@ class Config(object): self.logo = logo self.mongo = mongo self.urls = urls + self.cogs = cogs self.max_messages = max_messages self.gitlab_token = gitlab_token - db = DBManager(config=mongo).mongo.jarvis.config + db = DBManager(config=mongo["connect"]).mongo[mongo["database"]].config db_config = db.find() for item in db_config: setattr(self, item["key"], item["value"]) diff --git a/jarvis/db/types.py b/jarvis/db/types.py index adc764b..2f0840c 100644 --- a/jarvis/db/types.py +++ b/jarvis/db/types.py @@ -4,9 +4,10 @@ from datetime import datetime from typing import Any, Optional from bson import ObjectId +from pymongo import ASCENDING, DESCENDING + from jarvis.config import get_config from jarvis.db import DBManager -from pymongo import ASCENDING, DESCENDING logger = logging.getLogger("mongodb") @@ -35,7 +36,9 @@ coll_lookup = { "Warning": "warns", } -db_instance = DBManager(get_config().mongo).mongo.jarvis +m_config = get_config().mongo + +db_instance = DBManager(m_config["connect"]).mongo[m_config["database"]] ################# diff --git a/jarvis/events/guild.py b/jarvis/events/guild.py index 708ffae..eb5ccb9 100644 --- a/jarvis/events/guild.py +++ b/jarvis/events/guild.py @@ -13,8 +13,9 @@ class GuildEventHandler(object): async def on_guild_join(self, guild): general = find(lambda x: x.name == "general", guild.channels) if general and general.permissions_for(guild.me).send_messages: + user = self.bot.user await general.send( - "Allow me to introduce myself. I am J.A.R.V.I.S., a virtual " + f"Allow me to introduce myself. I am {user.mention}, 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." diff --git a/jarvis/utils/__init__.py b/jarvis/utils/__init__.py index 5461c17..949fdac 100644 --- a/jarvis/utils/__init__.py +++ b/jarvis/utils/__init__.py @@ -6,8 +6,8 @@ from discord import Color, Embed from discord.ext import commands import jarvis.cogs -import jarvis.config import jarvis.db +from jarvis.config import get_config __all__ = ["field", "db", "cachecog", "permissions"] @@ -40,7 +40,9 @@ def get_prefix(bot, message): def get_extensions(path=jarvis.cogs.__path__) -> list: - return ["jarvis.cogs.{}".format(x.name) for x in iter_modules(path)] + config = get_config() + vals = config.cogs or [x.name for x in iter_modules(path)] + return ["jarvis.cogs.{}".format(x) for x in vals] def parse_color_hex(hex: str) -> Color: diff --git a/jarvis/utils/permissions.py b/jarvis/utils/permissions.py index 2f4a9bd..58d5458 100644 --- a/jarvis/utils/permissions.py +++ b/jarvis/utils/permissions.py @@ -1,12 +1,12 @@ from discord.ext import commands -import jarvis +from jarvis.config import get_config def user_is_bot_admin(): def predicate(ctx): - if getattr(jarvis.config.get_config(), "admins", None): - return ctx.author.id in jarvis.config.get_config().admins + if getattr(get_config(), "admins", None): + return ctx.author.id in get_config().admins else: return False diff --git a/run.py b/run.py index 58701d8..b99f3cb 100755 --- a/run.py +++ b/run.py @@ -7,6 +7,7 @@ from time import sleep import git import jarvis +from jarvis.config import get_config def run(): @@ -93,7 +94,7 @@ Command List: if __name__ == "__main__": freeze_support() - config = jarvis.config.get_config() + config = get_config() pid_file = Value("i", 0) jarvis_process = Process(target=run, name="jarvis") logo = jarvis.logo.get_logo(config.logo)