async def test_pubsub_receiver_iter(create_redis, server, loop): sub = await create_redis(server.tcp_address, loop=loop) pub = await create_redis(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) async def coro(mpsc): lst = [] async for msg in mpsc.iter(): lst.append(msg) return lst tsk = asyncio.ensure_future(coro(mpsc), loop=loop) snd1, = await sub.subscribe(mpsc.channel('chan:1')) snd2, = await sub.subscribe(mpsc.channel('chan:2')) snd3, = await sub.psubscribe(mpsc.pattern('chan:*')) await pub.publish_json('chan:1', {'Hello': 'World'}) await pub.publish_json('chan:2', ['message']) mpsc.stop() await asyncio.sleep(0, loop=loop) assert await tsk == [ (snd1, b'{"Hello": "World"}'), (snd3, (b'chan:1', b'{"Hello": "World"}')), (snd2, b'["message"]'), (snd3, (b'chan:2', b'["message"]')), ] assert not mpsc.is_active
async def test_stopped(create_connection, server, loop): sub = await create_connection(server.tcp_address, loop=loop) pub = await create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) await sub.execute_pubsub('subscribe', mpsc.channel('channel:1')) assert mpsc.is_active mpsc.stop() with logs('aioredis', 'DEBUG') as cm: await pub.execute('publish', 'channel:1', b'Hello') await asyncio.sleep(0, loop=loop) assert len(cm.output) == 1 # Receiver must have 1 EndOfStream message warn_messaege = ( "WARNING:aioredis:Pub/Sub listener message after stop: " "sender: <_Sender name:b'channel:1', is_pattern:False, receiver:" "<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'") assert cm.output == [warn_messaege] # assert (await mpsc.get()) is None with pytest.raises(ChannelClosedError): await mpsc.get() res = await mpsc.wait_message() assert res is False
async def test_stopped(create_connection, server, caplog): sub = await create_connection(server.tcp_address) pub = await create_connection(server.tcp_address) mpsc = Receiver() await sub.execute_pubsub("subscribe", mpsc.channel("channel:1")) assert mpsc.is_active mpsc.stop() caplog.clear() with caplog.at_level("DEBUG", "aioredis"): await pub.execute("publish", "channel:1", b"Hello") await asyncio.sleep(0) assert len(caplog.record_tuples) == 1 # Receiver must have 1 EndOfStream message message = ( "Pub/Sub listener message after stop: " "sender: <_Sender name:b'channel:1', is_pattern:False, receiver:" "<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'") assert caplog.record_tuples == [ ("aioredis", logging.WARNING, message), ] # assert (await mpsc.get()) is None with pytest.raises(ChannelClosedError): await mpsc.get() res = await mpsc.wait_message() assert res is False
class Messager: def __init__(self, inbound, outbound, loop): self.loop = loop self.conn = None self.inbound = inbound self.outbound = outbound self.receiver = None self.task = None self.replies = dict() self.expected = set() async def initialize(self): self.conn = await aioredis.create_redis_pool(("localhost", 6379), encoding="utf-8", maxsize=2) self.receiver = Receiver(loop=self.loop) await self.conn.subscribe(self.receiver.channel(self.inbound)) self.task = self.loop.create_task(self.fetcher()) print( f'Redis connection established, listening on {self.inbound}, sending on {self.outbound}' ) # FIXME - propper logging async def terminate(self): # terminate channels and disconnect from redis self.conn.unsubscribe(self.inbound) self.task.cancel() self.receiver.stop() self.conn.close() await self.conn.wait_closed() async def fetcher(self): async for sender, message in self.receiver.iter(encoding='utf-8', decoder=json.loads): channel = sender.name.decode() if channel == self.inbound: uid = message["uid"] if uid not in self.expected: print("Unexpected message!") print(message) else: self.expected.remove(uid) self.replies[uid] = message async def get_reply(self, data): try: return (await asyncio.wait_for(self._get_reply(data), 10))["reply"] except TimeoutError: raise Redisception("No reply received from the bot!") async def _get_reply(self, data): uid = str(uuid.uuid4()) self.expected.add(uid) data["uid"] = uid await self.conn.publish_json(self.outbound, data) while uid not in self.replies: await asyncio.sleep(0.1) reply = self.replies[uid] del self.replies[uid] return reply
async def test_stopped(create_connection, server, loop): sub = await create_connection(server.tcp_address, loop=loop) pub = await create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) await sub.execute_pubsub('subscribe', mpsc.channel('channel:1')) assert mpsc.is_active mpsc.stop() with logs('aioredis', 'DEBUG') as cm: await pub.execute('publish', 'channel:1', b'Hello') await asyncio.sleep(0, loop=loop) assert len(cm.output) == 1 # Receiver must have 1 EndOfStream message warn_messaege = ( "WARNING:aioredis:Pub/Sub listener message after stop: " "sender: <_Sender name:b'channel:1', is_pattern:False, receiver:" "<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'" ) assert cm.output == [warn_messaege] # assert (await mpsc.get()) is None with pytest.raises(ChannelClosedError): await mpsc.get() res = await mpsc.wait_message() assert res is False
async def producer(self): channel_names = self.channel_names or [] channel_patterns = self.channel_patterns or [] if not channel_names and not channel_patterns: return mpsc = Receiver() if channel_names: channels = [mpsc.channel(c) for c in channel_names] await self.redis.subscribe(*channels) if channel_patterns: tasks = set() for p in channel_patterns: tasks.add(self.redis.psubscribe(mpsc.pattern(p))) if tasks: await asyncio.wait(tasks) try: await self.receiver_reader(mpsc) finally: if channel_names: await self.redis.unsubscribe(*channel_names) if channel_patterns: tasks = set() for p in channel_patterns: tasks.add(self.redis.punsubscribe(p)) await asyncio.wait(tasks) mpsc.stop()
async def main(): await redis.connect() mpsc = Receiver(loop=asyncio.get_event_loop()) asyncio.ensure_future(reader(mpsc)) await redis._redis.subscribe(mpsc.channel('channel:1')) while True: try: await asyncio.sleep(10) print('hearbeat', flush=True) except Exception as ex: print(ex, flush=True) mpsc.stop() await redis.disconnect()
def test_stopped(create_connection, server, loop): sub = yield from create_connection(server.tcp_address, loop=loop) pub = yield from create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) yield from sub.execute_pubsub('subscribe', mpsc.channel('channel:1')) assert mpsc.is_active mpsc.stop() with pytest.logs('aioredis', 'DEBUG') as cm: yield from pub.execute('publish', 'channel:1', b'Hello') yield from asyncio.sleep(0, loop=loop) assert len(cm.output) == 1 warn_messaege = ( "WARNING:aioredis:Pub/Sub listener message after stop: " "<_Sender name:b'channel:1', is_pattern:False, receiver:" "<Receiver is_active:False, senders:1, qsize:0>>, b'Hello'") assert cm.output == [warn_messaege] with pytest.raises(ChannelClosedError): yield from mpsc.get() res = yield from mpsc.wait_message() assert res is False
class DashLink(BaseCog): def __init__(self, bot): super().__init__(bot) bot.loop.create_task(self.init()) self.redis_link: aioredis.Redis = None self.receiver = Receiver(loop=bot.loop) self.handlers = dict(question=self.question, update=self.update, user_guilds=self.user_guilds, user_guilds_end=self.user_guilds_end, guild_info_watch=self.guild_info_watch, guild_info_watch_end=self.guild_info_watch_end) self.question_handlers = dict( heartbeat=self.still_spinning, user_info=self.user_info_request, get_guild_settings=self.get_guild_settings, save_guild_settings=self.save_guild_settings, replace_guild_settings=self.replace_guild_settings, setup_mute=self.setup_mute, cleanup_mute=self.cleanup_mute, cache_info=self.cache_info, guild_user_perms=self.guild_user_perms) # The last time we received a heartbeat, the current attempt number, how many times we have notified the owner self.last_dash_heartbeat = [time.time(), 0, 0] self.last_update = datetime.now() self.to_log = dict() self.update_message = None if Configuration.get_master_var( "TRANSLATIONS", dict(SOURCE="SITE", CHANNEL=0, KEY="", LOGIN="", WEBROOT=""))["SOURCE"] == 'CROWDIN': self.handlers["crowdin_webhook"] = self.crowdin_webhook self.task = self._receiver() def cog_unload(self): self.bot.loop.create_task(self._unload()) async def _unload(self): for c in self.receiver.channels.values(): self.redis_link.unsubscribe(c) self.receiver.stop() self.redis_link.close() await self.redis_link.wait_closed() async def init(self): try: self.redis_link = await aioredis.create_redis_pool( (Configuration.get_master_var('REDIS_HOST', "localhost"), Configuration.get_master_var('REDIS_PORT', 6379)), encoding="utf-8", db=0, maxsize=2) # size 2: one send, one receive self.bot.loop.create_task(self._receiver()) if Configuration.get_master_var("DASH_OUTAGE")["outage_detection"]: self.bot.loop.create_task(self.dash_monitor()) await self.redis_link.subscribe( self.receiver.channel("dash-bot-messages")) await self.redis_link.publish_json( "bot-dash-messages", { 'type': 'cache_info', 'message': await self.cache_info() }) except OSError: await GearbotLogging.bot_log("Failed to connect to the dash!") async def dash_monitor(self): DASH_OUTAGE_INFO: dict = Configuration.get_master_var("DASH_OUTAGE") DASH_OUTAGE_CHANNEl = DASH_OUTAGE_INFO["dash_outage_channel"] MAX_WARNINGS = DASH_OUTAGE_INFO["max_bot_outage_warnings"] BOT_OUTAGE_PINGED_ROLES = DASH_OUTAGE_INFO["dash_outage_pinged_roles"] while True: if (time.time() - self.last_dash_heartbeat[0]) > 5: self.last_dash_heartbeat[1] += 1 if self.last_dash_heartbeat[ 1] >= 3 and self.last_dash_heartbeat[2] < MAX_WARNINGS: print( "The dashboard API keepalive hasn't responded in over 3 minutes!" ) self.last_dash_heartbeat[2] += 1 self.last_dash_heartbeat[1] = 0 if DASH_OUTAGE_CHANNEl: outage_message = DASH_OUTAGE_INFO["dash_outage_embed"] # Apply the timestamp outage_message["timestamp"] = datetime.now().isoformat( ) # Set the color to the format Discord understands outage_message["color"] = outage_message["color"] # Generate the custom message and role pings notify_message = DASH_OUTAGE_INFO[ "dash_outage_message"] if BOT_OUTAGE_PINGED_ROLES: pinged_roles = [] for role_id in BOT_OUTAGE_PINGED_ROLES: pinged_roles.append(f"<@&{role_id}>") notify_message += f" Pinging: {', '.join(pinged_roles)}" try: outage_channel = self.bot.get_channel( DASH_OUTAGE_CHANNEl) await outage_channel.send( notify_message, embed=Embed.from_dict(outage_message)) except Forbidden: GearbotLogging.error( "We couldn't access the specified channel, the notification will not be sent!" ) # Wait a little bit longer so the dashboard has a chance to update before we check await asyncio.sleep(65) async def _handle(self, sender, message): try: await self.handlers[message["type"]](message["message"]) except CancelledError: return except Exception as e: await TheRealGearBot.handle_exception("Dash message handling", self.bot, e, None, None, None, message) async def send_to_dash(self, channel, **kwargs): await self.redis_link.publish_json("bot-dash-messages", dict(type=channel, message=kwargs)) async def question(self, message): try: reply = dict(reply=await self.question_handlers[message["type"]] (message["data"]), state="OK", uid=message["uid"]) except UnauthorizedException: reply = dict(uid=message["uid"], state="Unauthorized") except ValidationException as ex: reply = dict(uid=message["uid"], state="Bad Request", errors=ex.errors) except CancelledError: return except Exception as ex: reply = dict(uid=message["uid"], state="Failed") await self.send_to_dash("reply", **reply) raise ex await self.send_to_dash("reply", **reply) async def _receiver(self): async for sender, message in self.receiver.iter(encoding='utf-8', decoder=json.loads): self.bot.loop.create_task(self._handle(sender, message)) async def still_spinning(self, _): self.last_dash_heartbeat[0] = time.time() self.last_dash_heartbeat[1] = 0 self.last_dash_heartbeat[2] = 0 return self.bot.latency async def user_info_request(self, message): user_id = message["user_id"] user_info = await self.bot.fetch_user(user_id) return_info = { "username": user_info.name, "discrim": user_info.discriminator, "avatar_url": str(user_info.avatar_url_as(size=256)), "bot_admin_status": await self.bot.is_owner(user_info) or user_id in Configuration.get_master_var("BOT_ADMINS", []) } return return_info async def user_guilds(self, message): user_id = int(message["user_id"]) self.bot.dash_guild_users.add(user_id) self.redis_link.publish_json( "bot-dash-messages", dict(type="guild_add", message=dict(user_id=user_id, guilds=DashUtils.get_user_guilds( self.bot, user_id)))) async def user_guilds_end(self, message): user_id = int(message["user_id"]) self.bot.dash_guild_users.remove(user_id) async def guild_user_perms(self, message): guild = self.bot.get_guild(int(message["guild_id"])) if guild is None: return 0 return DashUtils.get_guild_perms( guild.get_member(int(message["user_id"]))) @needs_perm(DASH_PERMS.ACCESS) async def guild_info_watch(self, message): # start tracking info guild_id, user_id = get_info(message) if guild_id not in self.bot.dash_guild_watchers: self.bot.dash_guild_watchers[guild_id] = set() self.bot.dash_guild_watchers[guild_id].add(user_id) await self.send_guild_info( self.bot.get_guild(guild_id).get_member(user_id)) async def guild_info_watch_end(self, message): guild_id, user_id = get_info(message) if guild_id in self.bot.dash_guild_watchers: users = self.bot.dash_guild_watchers[guild_id] users.remove(user_id) if len(users) is 0: del self.bot.dash_guild_watchers[guild_id] async def send_guild_info_update_to_all(self, guild): if guild.id in self.bot.dash_guild_watchers: for user in self.bot.dash_guild_watchers[guild.id]: await self.send_guild_info(guild.get_member(user)) async def send_guild_info(self, member): await self.send_to_dash("guild_update", user_id=member.id, guild_id=member.guild.id, info=DashUtils.assemble_guild_info( self.bot, member)) @needs_perm(DASH_PERMS.VIEW_CONFIG) async def get_guild_settings(self, message): section = Configuration.get_var(int(message["guild_id"]), message["section"]) section = { k: [str(rid) if isinstance(rid, int) else rid for rid in v] if isinstance(v, list) else str(v) if isinstance(v, int) and not isinstance(v, bool) else v for k, v in section.items() } return section @needs_perm(DASH_PERMS.ALTER_CONFIG) async def save_guild_settings(self, message): guild_id, user_id = get_info(message) guild = self.bot.get_guild(guild_id) return DashConfig.update_config_section(guild, message["section"], message["modified_values"], guild.get_member(user_id)) @needs_perm(DASH_PERMS.ALTER_CONFIG) async def replace_guild_settings(self, message): guild_id, user_id = get_info(message) guild = self.bot.get_guild(guild_id) return DashConfig.update_config_section(guild, message["section"], message["modified_values"], guild.get_member(user_id), replace=True) async def cache_info(self, message=None): return { 'languages': Translator.LANG_NAMES, 'logging': { k: list(v.keys()) for k, v in GearbotLogging.LOGGING_INFO.items() } } @needs_perm(DASH_PERMS.ALTER_CONFIG) async def setup_mute(self, message): await self.override_handler( message, "setup", dict(send_messages=False, add_reactions=False), dict(speak=False, connect=False, stream=False)) @needs_perm(DASH_PERMS.ALTER_CONFIG) async def cleanup_mute(self, message): await self.override_handler(message, "cleanup", None, None) async def override_handler(self, message, t, text, voice): guild = self.bot.get_guild(message["guild_id"]) if not DashConfig.is_numeric(message["role_id"]): raise ValidationException(dict(role_id="Not a valid id")) role = guild.get_role(int(message["role_id"])) if role is None: raise ValidationException(dict(role_id="Not a valid id")) if role.id == guild.id: raise ValidationException( dict( role_id="The @everyone role can't be used for muting people" )) if role.managed: raise ValidationException( dict( role_id= "Managed roles can not be assigned to users and thus won't work for muting people" )) user = await Utils.get_user(message["user_id"]) parts = { "role_name": Utils.escape_markdown(role.name), "role_id": role.id, "user": Utils.clean_user(user), "user_id": user.id } GearbotLogging.log_key(guild.id, f"config_mute_{t}_triggered", **parts) failed = [] for channel in guild.text_channels: try: if text is None: await channel.set_permissions(role, reason=Translator.translate( f'mute_{t}', guild.id), overwrite=None) else: await channel.set_permissions(role, reason=Translator.translate( f'mute_{t}', guild.id), **text) except Forbidden as ex: failed.append(channel.mention) for channel in guild.voice_channels: try: if voice is None: await channel.set_permissions(role, reason=Translator.translate( f'mute_{t}', guild.id), overwrite=None) else: await channel.set_permissions(role, reason=Translator.translate( f'mute_{t}', guild.id), **voice) except Forbidden as ex: failed.append( Translator.translate('voice_channel', guild.id, channel=channel.name)) await asyncio.sleep( 1 ) # delay logging so the channel overrides can get querried and logged GearbotLogging.log_key(guild.id, f"config_mute_{t}_complete", **parts) out = '\n'.join(failed) GearbotLogging.log_key( guild.id, f"config_mute_{t}_failed", **parts, count=len(failed), tag_on=None if len(failed) is 0 else f'```{out}```') # crowdin async def crowdin_webhook(self, message): code = message["info"]["language"] await Translator.update_lang(code) if (datetime.now() - self.last_update).seconds > 5 * 60: self.update_message = None self.to_log = dict() if code not in self.to_log: self.to_log[code] = 0 self.to_log[code] += 1 embed = Embed(color=Color(0x1183f6), timestamp=datetime.utcfromtimestamp(time.time()), description=f"**Live translation update summary!**\n" + '\n'.join(f"{Translator.LANG_NAMES[code]} : {count}" for code, count in self.to_log.items())) if self.update_message is None: self.update_message = await Translator.get_translator_log_channel( )(embed=embed) else: await self.update_message.edit(embed=embed) self.last_update = datetime.now() async def update(self, message): t = message["type"] if t == "update": await Update.update("whoever just pushed to master", self.bot) elif t == "upgrade": await Update.upgrade("whoever just pushed to master", self.bot) else: raise RuntimeError( "UNKNOWN UPDATE MESSAGE, IS SOMEONE MESSING WITH IT?") @commands.Cog.listener() async def on_guild_join(self, guild): for user in self.bot.dash_guild_users: member = guild.get_member(user) if member is not None: permission = DashUtils.get_guild_perms(member) if permission > 0: await self.send_to_dash("guild_add", user_id=user, guilds={ str(guild.id): { "id": str(guild.id), "name": guild.name, "permissions": permission, "icon": guild.icon } }) @commands.Cog.listener() async def on_guild_remove(self, guild): for user in self.bot.dash_guild_users: member = guild.get_member(user) if member is not None: permission = DashUtils.get_guild_perms(member) if permission > 0: await self.send_to_dash("guild_remove", user_id=user, guild=str(guild.id)) @commands.Cog.listener() async def on_guild_update(self, before, after): for user in self.bot.dash_guild_users: member = after.get_member(user) if member is not None: old = DashUtils.get_guild_perms(member) new = DashUtils.get_guild_perms(member) if old != new: await self._notify_user(member, old, new, after) elif before.name != after.name or before.icon != after.icon: await self._notify_user(member, 0, 15, after) @commands.Cog.listener() async def on_member_update(self, before, after): if after.id in self.bot.dash_guild_users: old = DashUtils.get_guild_perms(before) new = DashUtils.get_guild_perms(after) await self._notify_user(after, old, new, before.guild) @commands.Cog.listener() async def on_guild_role_update(self, before, after): for user in self.bot.dash_guild_users: member = after.guild.get_member(user) if member is not None and after in member.roles: new = DashUtils.get_guild_perms(member) await self._notify_user(member, 0 if new is not 0 else 15, new, after.guild) @commands.Cog.listener() async def _notify_user(self, user, old, new, guild): if old != new: if new is not 0: await self.send_to_dash("guild_add", user_id=user.id, guilds={ str(guild.id): { "id": str(guild.id), "name": guild.name, "permissions": new, "icon": guild.icon } }) if new is 0 and old is not 0: await self.send_to_dash("guild_remove", user_id=user.id, guild=str(guild.id))
from aioredis.pubsub import Receiver from aioredis.abc import AbcChannel mpsc = Receiver(loop=loop) async def reader(mpsc): async for channel, msg in mpsc.iter(): assert isinstance(channel, AbcChannel) print("Got {!r} in channel {!r}".format(msg, channel)) asyncio.ensure_future(reader(mpsc)) await redis.subscribe(mpsc.channel('channel:1'), mpsc.channel('channel:3'), mpsc.channel('channel:5')) await redis.psubscribe(mpsc.pattern('hello')) # publishing 'Hello world' into 'hello-channel' # will print this message: # when all is done: await redis.unsubscribe('channel:1', 'channel:3', 'channel:5') await redis.punsubscribe('hello') mpsc.stop() # any message received after stop() will be ignored.
class DashLink(BaseCog): def __init__(self, bot): super().__init__(bot) bot.loop.create_task(self.init()) self.redis_link = None self.receiver = Receiver(loop=bot.loop) self.handlers = dict( guild_perm_request=self.guild_perm_request ) self.recieve_handlers = dict( ) self.last_update = datetime.now() self.to_log = dict() self.update_message = None if Configuration.get_master_var("TRANSLATIONS", dict(SOURCE="SITE", CHANNEL=0, KEY= "", LOGIN="", WEBROOT=""))["SOURCE"] == 'CROWDIN': self.recieve_handlers["crowdin_webhook"] = self.crowdin_webhook self.task = self._receiver() def cog_unload(self): self.bot.loop.create_task(self._unload()) async def _unload(self): for c in self.receiver.channels.values(): self.redis_link.unsubscribe(c) self.receiver.stop() self.redis_link.close() await self.redis_link.wait_closed() async def init(self): try: self.redis_link = await aioredis.create_redis_pool( (Configuration.get_master_var('REDIS_HOST', "localhost"), Configuration.get_master_var('REDIS_PORT', 6379)), encoding="utf-8", db=0, maxsize=2) # size 2: one send, one receive self.bot.loop.create_task(self._receiver()) await self.redis_link.subscribe(self.receiver.channel("dash-bot-messages")) except OSError: await GearbotLogging.bot_log("Failed to connect to the dash!") async def _receiver(self): async for sender, message in self.receiver.iter(encoding='utf-8', decoder=json.loads): try: if message["type"] in self.recieve_handlers.keys(): await self.recieve_handlers[message["type"]](message) else: reply = dict(reply=await self.handlers[message["type"]](message), uid=message["uid"]) await self.redis_link.publish_json("bot-dash-messages", reply) except Exception as e: await TheRealGearBot.handle_exception("Dash message handling", self.bot, e, None, None, None, message) async def guild_perm_request(self, message): info = dict() for guid in message["guild_list"]: guid = int(guid) guild = self.bot.get_guild(guid) permission = 0 if guild is not None: member = guild.get_member(int(message["user_id"])) mod_roles = Configuration.get_var(guid, "MOD_ROLES") if member.guild_permissions.ban_members or any(r.id in mod_roles for r in member.roles): permission |= (1 << 0) # dash access permission |= (1 << 1) # infraction access admin_roles = Configuration.get_var(guid, "ADMIN_ROLES") if member.guild_permissions.administrator or any(r.id in admin_roles for r in member.roles): permission |= (1 << 0) # dash access permission |= (1 << 2) # config read access permission |= (1 << 3) # config write access if permission > 0: info[guid] = dict(name=guild.name, permissions=permission, icon=guild.icon_url_as(size=256)) return info #crowdin async def crowdin_webhook(self, message): code = message["info"]["language"] await Translator.update_lang(code) if (datetime.now() - self.last_update).seconds > 5*60: self.update_message = None self.to_log = dict() if code not in self.to_log: self.to_log[code] = 0 self.to_log[code] += 1 embed = Embed(color=Color(0x1183f6), timestamp=datetime.utcfromtimestamp(time.time()), description=f"**Live translation update summary!**\n" + '\n'.join( f"{Translator.LANG_NAMES[code]} : {count}" for code, count in self.to_log.items())) if self.update_message is None: self.update_message = await Translator.get_translator_log_channel()(embed=embed) else: await self.update_message.edit(embed=embed) self.last_update = datetime.now()
class RedisBackend(BroadcastBackend): def __init__(self, url: str): self.conn_url = url self._pub_conn: typing.Optional[aioredis.Redis] = None self._sub_conn: typing.Optional[aioredis.Redis] = None self._msg_queue: typing.Optional[asyncio.Queue] = None self._reader_task: typing.Optional[asyncio.Task] = None self._mpsc: typing.Optional[Receiver] = None async def connect(self) -> None: if self._pub_conn or self._sub_conn or self._msg_queue: logger.warning("connections are already setup but connect called again; not doing anything") return self._pub_conn = await aioredis.create_redis(self.conn_url) self._sub_conn = await aioredis.create_redis(self.conn_url) self._msg_queue = asyncio.Queue() # must be created here, to get proper event loop self._mpsc = Receiver() self._reader_task = asyncio.create_task(self._reader()) async def disconnect(self) -> None: if self._pub_conn and self._sub_conn: self._pub_conn.close() self._sub_conn.close() else: logger.warning("connections are not setup, invalid call to disconnect") self._pub_conn = None self._sub_conn = None self._msg_queue = None if self._mpsc: self._mpsc.stop() else: logger.warning("redis mpsc receiver is not set, cannot stop it") if self._reader_task: if self._reader_task.done(): self._reader_task.result() else: logger.debug("cancelling reader task") self._reader_task.cancel() self._reader_task = None async def subscribe(self, channel: str) -> None: if not self._sub_conn: logger.error(f"not connected, cannot subscribe to channel {channel!r}") return await self._sub_conn.subscribe(self._mpsc.channel(channel)) async def unsubscribe(self, channel: str) -> None: if not self._sub_conn: logger.error(f"not connected, cannot unsubscribe from channel {channel!r}") return await self._sub_conn.unsubscribe(channel) async def publish(self, channel: str, message: typing.Any) -> None: if not self._pub_conn: logger.error(f"not connected, cannot publish to channel {channel!r}") return await self._pub_conn.publish_json(channel, message) async def next_published(self) -> Event: if not self._msg_queue: raise RuntimeError("unable to get next_published event, RedisBackend is not connected") return await self._msg_queue.get() async def _reader(self) -> None: async for channel, msg in self._mpsc.iter(encoding="utf8", decoder=json.loads): if not isinstance(channel, AbcChannel): logger.error(f"invalid channel returned from Receiver().iter() - {channel!r}") continue channel_name = channel.name.decode("utf8") if not self._msg_queue: logger.error(f"unable to put new message from {channel_name} into queue, not connected") continue await self._msg_queue.put(Event(channel=channel_name, message=msg))
class DashLink: def __init__(self, bot) -> None: self.bot: GearBot = bot bot.loop.create_task(self.init()) self.redis_link = None self.receiver = Receiver(loop=bot.loop) self.handlers = dict(guild_perm_request=self.guild_perm_request) self.task = self._receiver() def __unload(self): self.bot.loop.create_task(self._unload()) async def _unload(self): for c in self.receiver.channels.values(): self.redis_link.unsubscribe(c) self.receiver.stop() self.redis_link.close() await self.redis_link.wait_closed() async def init(self): try: self.redis_link = await aioredis.create_redis_pool( (Configuration.get_master_var('REDIS_HOST', "localhost"), Configuration.get_master_var('REDIS_PORT', 6379)), encoding="utf-8", db=0, maxsize=2) # size 2: one send, one receive self.bot.loop.create_task(self._receiver()) await self.redis_link.subscribe( self.receiver.channel("dash-bot-messages")) except OSError: await GearbotLogging.bot_log("Failed to connect to the dash!") async def _receiver(self): async for sender, message in self.receiver.iter(encoding='utf-8', decoder=json.loads): try: reply = dict(reply=await self.handlers[message["type"]] (message), uid=message["uid"]) await self.redis_link.publish_json("bot-dash-messages", reply) except Exception as e: await TheRealGearBot.handle_exception("Dash message handling", self.bot, e, None, None, None, message) async def guild_perm_request(self, message): info = dict() for guid in message["guild_list"]: guid = int(guid) guild = self.bot.get_guild(guid) permission = 0 if guild is not None: member = guild.get_member(int(message["user_id"])) mod_roles = Configuration.get_var(guid, "MOD_ROLES") if member.guild_permissions.ban_members or any( r.id in mod_roles for r in member.roles): permission |= (1 << 0) # dash access permission |= (1 << 1) # infraction access admin_roles = Configuration.get_var(guid, "ADMIN_ROLES") if member.guild_permissions.administrator or any( r.id in admin_roles for r in member.roles): permission |= (1 << 0) # dash access permission |= (1 << 2) # config read access permission |= (1 << 3) # config write access if permission > 0: info[guid] = dict(name=guild.name, permissions=permission, icon=guild.icon_url_as(size=256)) return info