From ea7f14459d3babf1640de504903ee67cd3152878 Mon Sep 17 00:00:00 2001 From: Embolalia Date: Sun, 6 Dec 2015 13:02:21 -0500 Subject: [PATCH] Add tracking of users and their accounts This uses some of the code that @maxpowa wrote for #941, but gives a somewhat more intuitive API. It also paves the way for potentially adding direct support for away-notify and metadata-notify. --- sopel/bot.py | 3 + sopel/coretasks.py | 160 ++++++++++++++++++++++++++++++++++++------ sopel/tools/target.py | 64 +++++++++++++++++ 3 files changed, 206 insertions(+), 21 deletions(-) create mode 100644 sopel/tools/target.py diff --git a/sopel/bot.py b/sopel/bot.py index b4cc290076..210fbad087 100644 --- a/sopel/bot.py +++ b/sopel/bot.py @@ -99,6 +99,9 @@ def __init__(self, config, daemon=False): bitwise integer value, determined by combining the appropriate constants from `module`.""" + self.channels_ = dict() # name to chan obj + self.users = dict() # name to user obj + self.db = SopelDB(config) """The bot's database.""" diff --git a/sopel/coretasks.py b/sopel/coretasks.py index e6d2dec6a1..f33a343d60 100644 --- a/sopel/coretasks.py +++ b/sopel/coretasks.py @@ -16,12 +16,14 @@ from __future__ import unicode_literals, absolute_import, print_function, division +from random import randint import re import sys import time import sopel import sopel.module from sopel.tools import Identifier, iteritems +from sopel.tools.target import User, Channel import base64 from sopel.logger import get_logger @@ -31,6 +33,7 @@ LOGGER = get_logger(__name__) batched_caps = {} +who_reqs = {} # Keeps track of reqs coming from this module, rather than others def auth_after_register(bot): @@ -234,6 +237,11 @@ def track_nicks(bot, trigger): value = bot.privileges[channel].pop(old) bot.privileges[channel][new] = value + for channel in bot.channels_.values(): + channel.rename_user(old, new) + if old in bot.users: + bot.users[new] = bot.users.pop(old) + @sopel.module.rule('(.*)') @sopel.module.event('PART') @@ -241,14 +249,9 @@ def track_nicks(bot, trigger): @sopel.module.thread(False) @sopel.module.unblockable def track_part(bot, trigger): - if trigger.nick == bot.nick: - bot.channels.remove(trigger.sender) - del bot.privileges[trigger.sender] - else: - try: - del bot.privileges[trigger.sender][trigger.nick] - except KeyError: - pass + nick = trigger.nick + channel = trigger.sender + _remove_from_channel(bot, nick, channel) @sopel.module.rule('.*') @@ -258,17 +261,58 @@ def track_part(bot, trigger): @sopel.module.unblockable def track_kick(bot, trigger): nick = Identifier(trigger.args[1]) + channel = trigger.sender + _remove_from_channel(bot, nick, channel) + + +def _remove_from_channel(bot, nick, channel): if nick == bot.nick: - bot.channels.remove(trigger.sender) - del bot.privileges[trigger.sender] + del bot.privileges[channel] + bot.channels.remove(channel) + + if channel in bot.channels_: + del bot.channels_[channel] + + lost_users = [] + for nick_, user in bot.users.items(): + if channel in user.channels: + del user.channels[channel] + if not user.channels: + lost_users.append(nick_) + for nick_ in lost_users: + del bot.users[nick_] + else: + if nick in bot.privileges[channel]: + del bot.privileges[channel][nick] + + user = bot.users.get(nick) + if user and channel in user.channels: + bot.channels_[channel].clear_user(nick) + if not user.channels: + del bot.users[nick] + bot.channels_[channel] + + +def _accounts_enabled(bot): + return ('account-notify' in bot.enabled_capabilities and + 'extended-join' in bot.enabled_capabilities) + + +def _send_who(bot, channel): + if _accounts_enabled(bot): + # WHOX syntax, see http://faerion.sourceforge.net/doc/irc/whox.var + # Needed for accounts in who replies. The random integer is a param + # to identify the reply as one from this command, because if someone + # else sent it, we have no fucking way to know what the format is. + rand = str(randint(0, 999)) + while rand in who_reqs: + rand = str(randint(0, 999)) + who_reqs[rand] = channel + bot.write(['WHO', channel, 'a%nuacht,' + rand]) else: - # Temporary fix to stop KeyErrors from being sent to channel - # The privileges dict may not have all nicks stored at all times - # causing KeyErrors - try: - del bot.privileges[trigger.sender][nick] - except KeyError: - pass + # We might be on an old network, but we still care about keeping our + # user list updated + bot.write(['WHO', channel]) @sopel.module.rule('.*') @@ -280,8 +324,19 @@ def track_join(bot, trigger): if trigger.nick == bot.nick and trigger.sender not in bot.channels: bot.channels.append(trigger.sender) bot.privileges[trigger.sender] = dict() + if trigger.nick == bot.nick and trigger.sender not in bot.channels_: + bot.channels_[trigger.sender] = Channel(trigger.sender) + _send_who(bot, trigger.sender) bot.privileges[trigger.sender][trigger.nick] = 0 + user = bot.users.get(trigger.nick) + if user is None: + user = User(trigger.nick, trigger.user, trigger.host) + bot.channels_[trigger.sender].add_user(user) + + if len(trigger.args) > 1 and trigger.args[1] != '*' and _accounts_enabled(bot): + user.account = trigger.args[1] + @sopel.module.rule('.*') @sopel.module.event('QUIT') @@ -292,6 +347,9 @@ def track_quit(bot, trigger): for chanprivs in bot.privileges.values(): if trigger.nick in chanprivs: del chanprivs[trigger.nick] + for channel in bot.channels_.values(): + channel.clear_user(trigger.nick) + bot.users.pop(trigger.nick, None) @sopel.module.rule('.*') @@ -314,10 +372,11 @@ def recieve_cap_list(bot, trigger): if req[0] and req[2]: # Call it. req[2](bot, req[0] + trigger) - # Server is acknowledinge SASL for us. - elif (trigger.args[0] == bot.nick and trigger.args[1] == 'ACK' and - 'sasl' in trigger.args[2]): - recieve_cap_ack_sasl(bot) + # Server is acknowledging SASL for us. + elif trigger.args[1] == 'ACK': + if (trigger.args[0] == bot.nick and 'sasl' in trigger.args[2]): + recieve_cap_ack_sasl(bot) + bot.enabled_capabilities.add(trigger.args[2].strip()) def recieve_cap_ls_reply(bot, trigger): @@ -347,6 +406,15 @@ def recieve_cap_ls_reply(bot, trigger): # parse it, so we don't need to worry if it fails. bot._cap_reqs['multi-prefix'] = (['', 'coretasks', None, None],) + def acct_warn(bot, cap): + LOGGER.info('Server does not support {}, or it conflicts with a custom ' + 'module. User account validation is not available.'.format( + cap)) + if 'account-notify' not in bot._cap_reqs: + bot._cap_reqs['account-notify'] = (['', 'coretasks', None, acct_warn],) + if 'extended-join' not in bot._cap_reqs: + bot._cap_reqs['extended-join'] = (['', 'coretasks', None, acct_warn],) + for cap, reqs in iteritems(bot._cap_reqs): # At this point, we know mandatory and prohibited don't co-exist, but # we need to call back for optionals if they're also prohibited @@ -493,3 +561,53 @@ def blocks(bot, trigger): return else: bot.reply(STRINGS['huh']) + + +@sopel.module.event('ACCOUNT') +@sopel.module.rule('.*') +def account_notify(bot, trigger): + print(trigger.nick) + if trigger.nick not in bot.users: + bot.users[trigger.nick] = User(trigger.nick, trigger.user, trigger.host) + account = trigger.args[0] + if account == '*': + account = None + bot.users[trigger.nick].account = account + + +@sopel.module.event('354') +@sopel.module.rule('.*') +@sopel.module.priority('high') +@sopel.module.unblockable +def recv_whox(bot, trigger): + if (len(trigger.args) < 2 or trigger.args[1] not in who_reqs or + not _accounts_enabled(bot)): + # Ignored, some module probably called WHO + # TODO a separate 352 handler function + return + if len(trigger.args) != 7: + return LOGGER.warning('While populating `bot.accounts` a WHO response was malformed.') + print(trigger.args) + _, _, channel, user, host, nick, account = trigger.args + nick = Identifier(nick) + channel = Identifier(channel) + if nick not in bot.users: + bot.users[nick] = User(nick, user, host) + user = bot.users[nick] + if account == '0': + user.account = None + else: + user.account = account + if channel not in bot.channels_: + bot.channels_[channel] = Channel(channel) + bot.channels_[channel].add_user(user) + + +@sopel.module.event('315') +@sopel.module.rule('.*') +@sopel.module.priority('high') +@sopel.module.unblockable +def end_who(bot, trigger): + if not _accounts_enabled(bot): + return + who_reqs.pop(trigger.args[1], None) diff --git a/sopel/tools/target.py b/sopel/tools/target.py new file mode 100644 index 0000000000..63d14a82d9 --- /dev/null +++ b/sopel/tools/target.py @@ -0,0 +1,64 @@ +# coding=utf-8 +from __future__ import unicode_literals, absolute_import, print_function, division + +import functools +from sopel.tools import Identifier + + +@functools.total_ordering +class User(object): + def __init__(self, nick, user, host): + assert isinstance(nick, Identifier) + self.nick = nick + self.user = user + self.host = host + self.channels = {} + + hostmask = property(lambda self: '{}!{}@{}'.format(self.nick, self.user, + self.host)) + + def __eq__(self, other): + if not isinstance(other, User): + return NotImplemented + return self.name == other.name + + def __lt__(self, other): + if not isinstance(other, User): + return NotImplemented + return self.name < other.name + + +@functools.total_ordering +class Channel(object): + def __init__(self, name): + assert isinstance(name, Identifier) + self.name = name + self.users = {} + self.privileges = {} + + def clear_user(self, nick): + user = self.users[nick] + user.channels.pop(self.name, None) + del self.users[nick] + del self.privileges[nick] + + def add_user(self, user): + assert isinstance(user, User) + self.users[user.nick] = user + self.privileges[user.nick] = 0 + + def rename_user(self, old, new): + if old in self.users: + self.users[new] = self.users.pop(old) + if old in self.privileges: + self.privileges[new] = self.privileges.pop(old) + + def __eq__(self, other): + if not isinstance(other, Channel): + return NotImplemented + return self.name == other.name + + def __lt__(self, other): + if not isinstance(other, Channel): + return NotImplemented + return self.name < other.name