Merge pull request #82 from yannickgloster/wickerman-features

Added Features Wickerman Requested

bot.py

@@ -8,7 +8,7 @@
 from utils.csgo_server import CSGOServer
 
 
-__version__ = '1.4.0-beta'
+__version__ = '1.5.0-beta'
 __dev__ = 745000319942918303
 
 class Discord_10man(commands.Bot):
@@ -38,6 +38,9 @@ def __init__(self, config: dict, startup_extensions: List[str]):
         self.version: str = __version__
         self.queue_ctx: commands.Context = None
         self.queue_voice_channel: discord.VoiceChannel = None
+        self.match_size = 10
+        self.spectators: List[discord.Member] = []
+        self.connect_dm = False
 
         logger = logging.getLogger('discord')
         logger.setLevel(logging.DEBUG)

checks.py

@@ -9,7 +9,7 @@ async def voice_channel(ctx: commands.Context):
 
 
 async def ten_players(ctx: commands.Context):
-    if ctx.author.voice is not None and (len(ctx.author.voice.channel.members) < 10 and not ctx.bot.dev):
+    if ctx.author.voice is not None and (len(ctx.author.voice.channel.members) < ctx.bot.match_size and not ctx.bot.dev):
         raise commands.CommandError(message='There must be 10 members connected to the voice channel')
     return True
 

cogs/csgo.py

@@ -133,7 +133,7 @@ async def pug(self, ctx: commands.Context, *args):
             current_captain = team1_captain
             player_veto_count = 0
 
-            message = await ctx.send('10 man time\nLoading player selection...')
+            message = await ctx.send(f'{self.bot.match_size} man time\nLoading player selection...')
             for emoji in emojis:
                 await message.add_reaction(emoji)
 
@@ -225,6 +225,7 @@ async def pug(self, ctx: commands.Context, *args):
 
         team1_steamIDs = []
         team2_steamIDs = []
+        spectator_steamIDs = []
 
         if ctx.author.voice.channel.category is None:
             team1_channel = await ctx.guild.create_voice_channel(name=f'team_{team1_captain.display_name}',
@@ -249,6 +250,12 @@ async def pug(self, ctx: commands.Context, *args):
                                       {"player": str(player.id)})
             team2_steamIDs.append(data[0])
 
+        if len(self.bot.spectators) > 0:
+            for spec in self.bot.spectators:
+                data = await db.fetch_one('SELECT steam_id FROM users WHERE discord_id = :spectator',
+                                          {"spectator": str(spec.id)})
+                spectator_steamIDs.append(data[0])
+
         if map_arg is None:
             map_list = await self.map_veto(ctx, team1_captain, team2_captain)
         else:
@@ -309,8 +316,9 @@ async def pug(self, ctx: commands.Context, *args):
             'skip_veto': True,
             'veto_first': 'team1',
             'side_type': 'always_knife',
-            'players_per_team': len(team2),
+            'players_per_team': self.bot.match_size,
             'min_players_to_ready': 1,
+            'spectators': spectator_steamIDs,
             'team1': {
                 'name': f'team_{team1_captain.display_name}',
                 'tag': 'team1',
@@ -344,7 +352,11 @@ async def pug(self, ctx: commands.Context, *args):
 
         await asyncio.sleep(5)
         connect_embed = await self.connect_embed(csgo_server)
-        await ctx.send(embed=connect_embed)
+        if self.bot.connect_dm:
+            for player in team1 + team2 + self.bot.spectators:
+                await player.send(embed=connect_embed)
+        else:
+            await ctx.send(embed=connect_embed)
         score_embed = discord.Embed()
         score_embed.add_field(name='0', value=f'team_{team1_captain.display_name}', inline=True)
         score_embed.add_field(name='0', value=f'team_{team2_captain.display_name}', inline=True)
@@ -535,7 +547,7 @@ async def queue_check(self):
             if server.available:
                 available = True
                 break
-        if (len(self.bot.queue_voice_channel.members) >= 10 or (self.bot.dev and len(self.bot.queue_voice_channel.members) >= 1)) and available:
+        if (len(self.bot.queue_voice_channel.members) >= self.bot.match_size or (self.bot.dev and len(self.bot.queue_voice_channel.members) >= 1)) and available:
             embed = discord.Embed()
             embed.add_field(name='You have 60 seconds to ready up!', value='Ready: ✅', inline=False)
             ready_up_message = await self.bot.queue_ctx.send(embed=embed)
@@ -548,7 +560,7 @@ async def queue_check(self):
     async def ready_up(self, message: discord.Message, members: List[discord.Member]):
         message = await self.bot.queue_ctx.fetch_message(message.id)
 
-        # TODO: Add check for only the first 10 users
+        # TODO: Add check for only the first self.bot.match_size users
         check_emoji = None
         for reaction in message.reactions:
             if reaction.emoji == '✅':

cogs/setup.py

@@ -39,6 +39,57 @@ async def link_error(self, ctx: commands.Context, error: Exception):
             await ctx.send(str(error))
         traceback.print_exc()
 
+    @commands.command(aliases=['spectator', 'spec'],
+                      help='Adds this user as a spectator in the config for the next map.',
+                      brief='Add user as spectator', usage='<@User>')
+    @commands.has_permissions(administrator=True)
+    async def add_spectator(self, ctx: commands.Context, spec: discord.Member):
+        db = Database('sqlite:///main.sqlite')
+        await db.connect()
+        data = await db.fetch_one('SELECT 1 FROM users WHERE discord_id = :spectator', {"spectator": str(spec.id)})
+        if data is None:
+            raise commands.UserInputError(message=f'<@{spec.id}> needs to `.link` their account.')
+        self.bot.spectators.append(spec)
+        await ctx.send(f'<@{spec.id}> was added as a spectator.')
+
+    @add_spectator.error
+    async def add_spectator_error(self, ctx: commands.Context, error: Exception):
+        if isinstance(error, commands.UserInputError):
+            await ctx.send(str(error))
+        traceback.print_exc()
+
+    @commands.command(aliases=['remove_spec'],
+                      help='Removes this user as a spectator from the config.',
+                      brief='Removes user as spectator', usage='<@User>')
+    @commands.has_permissions(administrator=True)
+    async def remove_spectator(self, ctx: commands.Context, spec: discord.Member):
+        db = Database('sqlite:///main.sqlite')
+        await db.connect()
+        data = await db.fetch_one('SELECT 1 FROM users WHERE discord_id = :spectator',
+                                  {"spectator": str(spec.id)})
+        if data is None:
+            raise commands.UserInputError(message=f'User did not `.link` their account and probably is not a spectator.')
+        if data[0] in self.bot.spectators:
+            self.bot.spectators.remove(spec)
+            await ctx.send(f'<@{spec.id}> was added as a spectator.')
+        else:
+            raise commands.CommandError(message=f'<@{spec.id}> is not a spectator.')
+
+    @remove_spectator.error
+    async def remove_spectator_error(self, ctx: commands.Context, error: Exception):
+        if isinstance(error, commands.UserInputError) or isinstance(error, commands.CommandError):
+            await ctx.send(str(error))
+        traceback.print_exc()
+
+    @commands.command(aliases=['dm'],
+                      help='Command to enable or disable sending a dm with the connect ip vs posting it in the channel',
+                      brief='Enable or disable connect via dm')
+    @commands.check(checks.voice_channel)
+    @commands.has_permissions(administrator=True)
+    async def connect_dm(self, ctx: commands.Context, enabled: bool = False):
+        self.bot.connect_dm = enabled
+        await ctx.send(f'Connect message will {"not" if not enabled else ""} be sent via a DM.')
+
     @commands.command(aliases=['setupqueue'],
                       help='Command to set the server for the queue system. You must be in a voice channel.',
                       brief='Set\'s the server for the queue')
@@ -66,6 +117,27 @@ async def setup_queue_error(self, ctx: commands.Context, error: Exception):
             await ctx.send(str(error))
         traceback.print_exc()
 
+    @commands.command(aliases=['setup_queue_size', 'match_size', 'queue_size', 'set_match_size', 'set_queue_size'],
+                      help='This command sets the size of the match and the queue.',
+                      brief='Sets the size of the match & queue', usage='<size>')
+    @commands.has_permissions(administrator=True)
+    async def setup_match_size(self, ctx: commands.Context, size: int):
+        if size <= 0:
+            raise commands.CommandError(message=f'Invalid match size.')
+        if size % 2 != 0:
+            raise commands.CommandError(message=f'Match size must be an even number.')
+        self.bot.match_size = size
+        await ctx.send(f'Set match size to {self.bot.match_size}.')
+
+    @setup_match_size.error
+    async def setup_match_size_error(self, ctx: commands.Context, error: Exception):
+        if isinstance(error, commands.BadArgument):
+            await ctx.send('Invalid Argument')
+        elif isinstance(error, commands.CommandError):
+            await ctx.send(str(error))
+        traceback.print_exc()
+
+
     @commands.command(help='Command to send a test message to the server to verify that RCON is working.',
                       brief='Sends a message to the server to test RCON', usage='<message>')
     @commands.has_permissions(administrator=True)
@@ -83,7 +155,7 @@ async def RCON_message_error(self, ctx: commands.Context, error: Exception):
             await ctx.send('Please specify the message')
         traceback.print_exc()
 
-    @commands.command(help='This command unbans everyone on the server. Useful fix',
+    @commands.command(help='This command unbans everyone on the server. Useful fix.',
                       brief='Unbans everyone from the server', hidden=True)
     @commands.has_permissions(administrator=True)
     async def RCON_unban(self, ctx: commands.Context):
@@ -98,6 +170,13 @@ async def RCON_unban_error(self, ctx: commands.Context, error: Exception):
             await ctx.send('Only an administrator can unban every player')
         traceback.print_exc()
 
+    @commands.command(aliases=['end', 'stop'],
+                      help='This command force ends a match.',
+                      brief='Force ends a match', usage='<ServerID>')
+    @commands.has_permissions(administrator=True)
+    async def force_end(self, ctx: commands.Context, server_id: int = 0):
+        valve.rcon.execute((self.bot.servers[server_id].server_address, self.bot.servers[server_id].server_port),
+                           self.bot.servers[server_id].RCON_password, 'get5_endmatch')
 
 def setup(client):
     client.add_cog(Setup(client))