You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

93 lines
3.8 KiB

import re
import time
import discord
from commands import CommandError
def cooldown(retries, timeout, failure):
def do_cooldown(function):
def wrapper(self, user, *args, **kwargs):
cooldowns = getattr(function, 'cooldowns', {})
if user not in cooldowns:
cooldowns[user] = dict(tries=0, last_run=0)
cd = cooldowns[user]
if cd['tries'] < retries:
cd['tries'] += 1
cd['last_run'] = time.time()
setattr(function, 'cooldowns', cooldowns)
return function(self, user, *args, **kwargs)
if time.time() - cd['last_run'] > timeout:
cd['tries'] = 1
cd['last_run'] = time.time()
setattr(function, 'cooldowns', cooldowns)
return function(self, user, *args, **kwargs)
return failure(self, user, *args, **kwargs)
return wrapper
return do_cooldown
class DiscordClient(discord.Client):
def __init__(self, config, logger, commands):
self.config = config
self.logger = logger
self.commands = commands
self.supported_commands = [
(re.compile(r'^(?P<prefix>!|\?)(bella(gram|pics)|insta(gram|bella))$'), self._do_bellagram),
(re.compile(r'^(?P<prefix>!|\?)yt\s+(?P<q>")?(?P<query>.+)(?(q)")$'), self._do_yt),
]
super(DiscordClient, self).__init__()
async def start_(self):
token = self.config['Discord'].get('token')
await self.start(token)
async def on_ready(self):
self.logger.info('Logged in as {0}'.format(self.user.name))
async def on_message(self, message):
for pattern, action in self.supported_commands:
m = pattern.match(message.content)
if m:
await action(message.author.id, message, **m.groupdict())
async def _cooldown_failure(self, user, message, **kwargs):
await self.send_message(message.channel,
'Sorry {0}, you have to wait a while before running '
'the same command again'.format(message.author.mention))
@cooldown(retries=2, timeout=5*60, failure=_cooldown_failure)
async def _do_bellagram(self, user, message, prefix, **kwargs):
try:
bellagram = self.commands.bellagram()
except CommandError as e:
await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e))
else:
if prefix == '!':
embed = discord.Embed(title=bellagram['title'], url=bellagram['url'], color=0x8545bc)
embed.set_image(url=bellagram['display_url'])
embed.set_author(name=bellagram['owner'], url=bellagram['owner_url'],
icon_url=bellagram['owner_pic_url'])
await self.send_message(message.channel, embed=embed)
elif prefix == '?':
await self.send_message(message.channel, bellagram['url'])
@cooldown(retries=3, timeout=5*60, failure=_cooldown_failure)
async def _do_yt(self, user, message, query, prefix, **kwargs):
try:
result = self.commands.query_youtube(query)
except CommandError as e:
await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e))
else:
if prefix == '!':
embed = discord.Embed(title=result['title'], url=result['url'],
description=result['description'], color=0xff0000)
embed.set_thumbnail(url=result['thumbnail_url'])
embed.set_author(name=result['channel_title'], url=result['channel_url'],
icon_url=result['channel_thumbnail_url'])
await self.send_message(message.channel, embed=embed)
elif prefix == '?':
await self.send_message(message.channel, result['url'])