Add whitelist command

Author: Adikso

README.md

@@ -82,6 +82,7 @@ Command | Alias | Usage | Description
 .macro | .m | .macro <macro_name> | Executes macro from "macros" directory
 .interactive | .i | .interactive | Enables interactive mode where all messages are sent to terminal
 .exec | .e | .exec <command...> | Execute single command
+.whitelist | - | .whitelist add/remove <user_tag> | Add/Remove user to/from whitelist. **Bot owner can execute all commands without being on the whitelist**
 
 ### Shortcuts
 Shortcut | Description

bashbot/bot.py

@@ -17,6 +17,7 @@
 from bashbot.command.select import SelectCommand
 from bashbot.command.submit import SubmitCommand
 from bashbot.command.help import HelpCommand
+from bashbot.command.whitelist import WhitelistCommand
 from bashbot.core.exceptions import SessionDontExistException, ArgumentFormatException, TerminalNotFoundException, \
     MacroNotFoundException
 from bashbot.core.settings import settings
@@ -46,6 +47,7 @@ def __init__(self, command_prefix, **options):
         self.add_cog(InteractiveCommand())
         self.add_cog(SubmitCommand())
         self.add_cog(ExecCommand())
+        self.add_cog(WhitelistCommand())
 
         self.remove_command("help")
         self.add_cog(HelpCommand())
@@ -81,10 +83,27 @@ async def on_message(self, message: Message):
         if message.author.bot:
             return
 
-        if isinstance(message.channel, DMChannel) and settings().get('discord.disable_dm'):
-            embed = Embed(title=f'Using bot on DM is disabled', description='discord.disable_dm = true')
-            await message.channel.send(embed=embed)
-            return
+        is_owner = await self.is_owner(message.author)
+        if not is_owner:
+            if settings().get('discord.enable_users_whitelist'):
+                users_whitelist = settings().get('discord.users_whitelist', [])
+
+                if message.author.id not in users_whitelist:
+                    first_prefix = settings().get('commands.prefixes')[0]
+                    embed = Embed(
+                        title=f'Only whitelisted users can execute commands',
+                        description=f'{first_prefix}.whitelist add {message.author.mention}'
+                    )
+                    await message.channel.send(embed=embed)
+                    return
+
+            if isinstance(message.channel, DMChannel) and settings().get('discord.disable_dm'):
+                embed = Embed(
+                    title=f'Using bot on DM is disabled',
+                    description='discord.disable_dm = true'
+                )
+                await message.channel.send(embed=embed)
+                return
 
         terminal = sessions().by_channel(message.channel)
 

bashbot/command/__init__.py

@@ -13,9 +13,3 @@ async def predicate(ctx):
 
     return commands.check(predicate)
 
-
-def has_permission(name):
-    async def predicate(ctx):
-        return True
-
-    return commands.check(predicate)

bashbot/command/about.py

@@ -2,7 +2,6 @@
 from discord.ext import commands
 from discord.ext.commands import Context
 
-from bashbot.command import has_permission
 from bashbot.constants import REPOSITORY_URL, REPOSITORY_AUTHOR, THUMBNAIL_URL, EMBED_COLOR
 from bashbot.core.settings import settings
 from bashbot.core.updater import updater
@@ -13,7 +12,6 @@ class AboutCommand(commands.Cog):
         name='.about',
         description='Shows information about project'
     )
-    @has_permission('info.about')
     async def about(self, ctx: Context):
         embed = Embed(title='About BashBot', description='BashBot is a Discord bot that allows terminal access via chat.', color=EMBED_COLOR)
         embed.add_field(name='Github', value=REPOSITORY_URL, inline=False)

bashbot/command/close.py

@@ -1,6 +1,6 @@
 from discord.ext import commands
 
-from bashbot.command import has_permission, session_exists
+from bashbot.command import session_exists
 from bashbot.core.exceptions import SessionDontExistException
 from bashbot.core.settings import settings
 from bashbot.terminal.sessions import sessions
@@ -12,7 +12,6 @@ class CloseCommand(commands.Cog):
         aliases=['.c'],
         description='Closes current terminal session'
     )
-    @has_permission('session.close')
     @session_exists()
     async def close(self, ctx):
         terminal = sessions().by_channel(ctx.message.channel)

bashbot/command/controls.py

@@ -1,7 +1,7 @@
 from discord import Message
 from discord.ext import commands
 
-from bashbot.command import session_exists, has_permission
+from bashbot.command import session_exists
 from bashbot.terminal.sessions import sessions
 from bashbot.terminal.terminal import Terminal
 
@@ -18,7 +18,6 @@ async def controls(self, ctx):
 
     @controls.command()
     @session_exists()
-    @has_permission('input.controls.manage')
     async def add(self, ctx, emoji_id, content):
         terminal: Terminal = sessions().by_channel(ctx.channel)
         message: Message = sessions().find_message(terminal)
@@ -28,7 +27,6 @@ async def add(self, ctx, emoji_id, content):
 
     @controls.command()
     @session_exists()
-    @has_permission('input.controls.manage')
     async def remove(self, ctx, emoji_id):
         terminal: Terminal = sessions().by_channel(ctx.channel)
         message: Message = sessions().find_message(terminal)

bashbot/command/exec.py

@@ -4,7 +4,6 @@
 
 from discord.ext import commands
 
-from bashbot.command import has_permission
 from bashbot.core.settings import settings
 
 
@@ -15,7 +14,6 @@ class ExecCommand(commands.Cog):
         description='Execute single command',
         usage='<command...>'
     )
-    @has_permission('session.open')
     async def exec(self, ctx, *, command):
         shell_path = settings().get('terminal.shell_path')
 

bashbot/command/freeze.py

@@ -1,6 +1,6 @@
 from discord.ext import commands
 
-from bashbot.command import has_permission, session_exists
+from bashbot.command import session_exists
 from bashbot.core.exceptions import SessionDontExistException
 from bashbot.terminal.sessions import sessions
 from bashbot.terminal.terminal import TerminalState
@@ -12,7 +12,6 @@ class FreezeCommand(commands.Cog):
         aliases=['.f'],
         description='Freezes current terminal session'
     )
-    @has_permission('session.freeze')
     @session_exists()
     async def freeze(self, ctx):
         terminal = sessions().by_channel(ctx.message.channel)

bashbot/command/help.py

@@ -2,7 +2,6 @@
 from discord.ext import commands
 from discord.ext.commands import Context
 
-from bashbot.command import has_permission
 from bashbot.constants import THUMBNAIL_URL, EMBED_COLOR, EMBED_COLOR_ERROR
 from bashbot.core.settings import settings
 
@@ -12,7 +11,6 @@ class HelpCommand(commands.Cog):
         name='.help',
         description='This command'
     )
-    @has_permission('help')
     async def help(self, ctx: Context, command_name=None):
         if command_name:
             await self._help_for_command(ctx, command_name)

bashbot/command/here.py

@@ -1,7 +1,7 @@
 from discord.ext import commands
 from discord.ext.commands import Context
 
-from bashbot.command import has_permission, session_exists
+from bashbot.command import session_exists
 from bashbot.terminal.sessions import sessions
 
 
@@ -11,7 +11,6 @@ class HereCommand(commands.Cog):
         aliases=['.h'],
         description='Moves selected terminal below the user message'
     )
-    @has_permission('session.here')
     @session_exists()
     async def here(self, ctx: Context):
         terminal = sessions().by_channel(ctx.channel)

bashbot/command/interactive.py

@@ -1,7 +1,7 @@
 from discord.ext import commands
 from discord.ext.commands import Context
 
-from bashbot.command import has_permission, session_exists
+from bashbot.command import session_exists
 from bashbot.core.exceptions import SessionDontExistException
 from bashbot.terminal.sessions import sessions
 
@@ -12,7 +12,6 @@ class InteractiveCommand(commands.Cog):
         aliases=['.i'],
         description='Toggles interactive mode where all messages are sent to terminal'
     )
-    @has_permission('session.interactive')
     @session_exists()
     async def interactive(self, ctx: Context):
         terminal = sessions().by_channel(ctx.message.channel)

bashbot/command/macro.py

@@ -1,7 +1,7 @@
 from discord.ext import commands
 from discord.ext.commands import Context
 
-from bashbot.command import has_permission, session_exists
+from bashbot.command import session_exists
 from bashbot.core.macros import execute_macro
 
 
@@ -12,7 +12,6 @@ class MacroCommand(commands.Cog):
         description='Executes macro from "macros" directory',
         usage='<macro_name>'
     )
-    @has_permission('session.macro')
     @session_exists()
     async def macro(self, ctx: Context, macro_name: str):
         await execute_macro(ctx, macro_name)

bashbot/command/open.py

@@ -1,6 +1,5 @@
 from discord.ext import commands
 
-from bashbot.command import has_permission
 from bashbot.core.exceptions import ArgumentFormatException
 from bashbot.core.macros import execute_macro
 from bashbot.core.settings import settings
@@ -16,7 +15,6 @@ class OpenCommand(commands.Cog):
         description='Opens new terminal session',
         usage='[name]'
     )
-    @has_permission('session.open')
     async def open(self, ctx, name: str = None):
         if name and len(name) > 20:
             raise ArgumentFormatException('Session name length exceeds 20 characters limit')

bashbot/command/rename.py

@@ -1,6 +1,6 @@
 from discord.ext import commands
 
-from bashbot.command import has_permission, session_exists
+from bashbot.command import session_exists
 from bashbot.terminal.sessions import sessions
 
 
@@ -10,7 +10,6 @@ class RenameCommand(commands.Cog):
         description='Changes session name',
         usage='<new_name>'
     )
-    @has_permission('session.rename')
     @session_exists()
     async def rename(self, ctx, new_name):
         terminal = sessions().by_channel(ctx.message.channel)

bashbot/command/repeat.py

@@ -1,6 +1,5 @@
 from discord.ext import commands
 
-from bashbot.command import has_permission
 from bashbot.core.exceptions import SessionDontExistException
 from bashbot.terminal.sessions import sessions
 
@@ -12,7 +11,6 @@ class RepeatCommand(commands.Cog):
         description='Repeats string n times and sends to the current terminal session',
         usage='<string..>'
     )
-    @has_permission('session.repeat')
     async def repeat(self, ctx, n: int, *, text):
         terminal = sessions().by_channel(ctx.message.channel)
         if not terminal:

bashbot/command/select.py

@@ -1,6 +1,5 @@
 from discord.ext import commands
 
-from bashbot.command import has_permission
 from bashbot.core.exceptions import TerminalNotFoundException
 from bashbot.terminal.sessions import sessions
 
@@ -12,7 +11,6 @@ class SelectCommand(commands.Cog):
         description='Sets terminal as selected',
         usage='[name]'
     )
-    @has_permission('session.select')
     async def select(self, ctx, name):
         terminal = sessions().by_name(name)
         if not terminal:

bashbot/command/submit.py

@@ -1,7 +1,7 @@
 from discord.ext import commands
 from discord.ext.commands import Context
 
-from bashbot.command import has_permission, session_exists
+from bashbot.command import session_exists
 from bashbot.terminal.sessions import sessions
 
 
@@ -10,7 +10,6 @@ class SubmitCommand(commands.Cog):
         name='.submit',
         description='Toggles auto submit mode'
     )
-    @has_permission('session.submit')
     @session_exists()
     async def submit(self, ctx: Context):
         terminal = sessions().by_channel(ctx.channel)

bashbot/command/whitelist.py

@@ -0,0 +1,35 @@
+from discord import User
+from discord.ext import commands
+
+from bashbot.core.settings import settings
+
+
+class WhitelistCommand(commands.Cog):
+    @commands.group(
+        name='.whitelist',
+        description='Manages users whitelist',
+        usage='add/remove <user_tag>'
+    )
+    async def controls(self, ctx):
+        if ctx.invoked_subcommand is None:
+            pass
+
+    @controls.command()
+    async def add(self, ctx, user: User):
+        whitelist = settings().config['discord']['users_whitelist']
+        if user.id not in whitelist:
+            whitelist.append(user.id)
+            settings().save()
+            await ctx.send(f"Added user {user.mention} to whitelist")
+        else:
+            await ctx.send(f"User {user.mention} is already on the whitelist")
+
+    @controls.command()
+    async def remove(self, ctx, user: User):
+        whitelist = settings().config['discord']['users_whitelist']
+        if user.id in whitelist:
+            whitelist.remove(user.id)
+            settings().save()
+            await ctx.send(f"Removed user {user.mention} from whitelist")
+        else:
+            await ctx.send(f"User {user.mention} is not on the whitelist")

bashbot/core/settings.py

@@ -23,6 +23,8 @@ def load(self, path=DEFAULT_CONFIG_PATH):
         self.add_default('discord.token', 'TOKEN_HERE')
         self.add_default('discord.presence', '{prefix}.help')
         self.add_default('discord.disable_dm', True)
+        self.add_default('discord.enable_users_whitelist', True)
+        self.add_default('discord.users_whitelist', [])
 
         # [terminal]
         self.add_default('terminal.template', '`| TTY #{name} | {state} |`\n```{content}```')