diff --git a/bot.py b/bot.py index 23d68ea..b8f70d1 100644 --- a/bot.py +++ b/bot.py @@ -1,7 +1,9 @@ import inspect +import itertools import logging import os +import discord from discord.ext import commands import lib @@ -9,6 +11,7 @@ import heavynode DISCORD_TOKEN = os.environ['discord_token'] +DISCORD_SERVER_ID = 530446700058509323 HEAVYNODE_TOKEN = os.environ['heavynode_token'] @@ -19,8 +22,19 @@ hn = heavynode.Client(HEAVYNODE_TOKEN) bot.add_cleanup(hn.shutdown) +async def is_admin(ctx): + user = ctx.message.author + guild = discord.utils.get(bot.guilds, id=DISCORD_SERVER_ID) + member = discord.utils.get(guild.members, id=user.id) + if member is not None: + for role in member.roles: + if role.name == 'Admin' or role.name == 'Mod': + return True + return False + + @bot.command() -@commands.has_any_role('Admin', 'Mod') +@commands.check(is_admin) async def add(ctx, player): """Add a player to the server whitelist. Must use exact Minecraft name.""" await hn.send_command(f'whitelist add {player}') @@ -28,7 +42,7 @@ async def add(ctx, player): @bot.command() -@commands.has_any_role('Admin', 'Mod') +@commands.check(is_admin) async def remove(ctx, player): """Remove a player from the server whitelist. Must use exact Minecraft name.""" await hn.send_command(f'whitelist remove {player}')