async def sasl_ack(conn): sasl_auth = conn.config.get('sasl') if sasl_auth and sasl_auth.get('enabled', True): sasl_mech = sasl_auth.get("mechanism", "PLAIN").upper() auth_fut = async_util.create_future(conn.loop) conn.memory["sasl_auth_future"] = auth_fut conn.cmd("AUTHENTICATE", sasl_mech) cmd, arg = await auth_fut if cmd == "908": logger.warning("[%s|sasl] SASL mechanism not supported", conn.name) elif cmd == "AUTHENTICATE" and arg[0] == '+': num_fut = async_util.create_future(conn.loop) conn.memory["sasl_numeric_future"] = num_fut if sasl_mech == "PLAIN": auth_str = "{user}\0{user}\0{passwd}".format( user=sasl_auth["user"], passwd=sasl_auth["pass"]).encode() conn.cmd("AUTHENTICATE", base64.b64encode(auth_str).decode()) else: conn.cmd("AUTHENTICATE", '+') numeric = await num_fut if numeric == "902": logger.warning("[%s|sasl] SASL nick locked", conn.name) elif numeric == "903": logger.info("[%s|sasl] SASL auth successful", conn.name) elif numeric == "904": logger.warning("[%s|sasl] SASL auth failed", conn.name) elif numeric == "905": logger.warning("[%s|sasl] SASL auth too long", conn.name) elif numeric == "906": logger.warning("[%s|sasl] SASL auth aborted", conn.name) elif numeric == "907": logger.warning("[%s|sasl] SASL auth already completed", conn.name)
def get_initial_connection_data(conn, loop, db, event): """ - Update all user data :type conn: cloudbot.client.Client :type loop: asyncio.AbstractEventLoop :type db: sqlalchemy.orm.Session :type event: cloudbot.event.Event """ if conn.nick.endswith('-dev') and not hasattr(event, 'triggered_command'): # Ignore initial data update on development instances return fut = create_future(loop) try: lines, fut = conn.memory["sherlock"]["futures"]["who_0"][0] except LookupError: pass else: if not fut.done(): return "getdata command already running" conn.memory["sherlock"]["futures"]["who_0"][0] = ([], fut) yield from asyncio.sleep(10) now = datetime.datetime.now() conn.send("WHO 0") try: lines = yield from asyncio.wait_for(fut, 30 * 60) except asyncio.TimeoutError: return "Timeout reached" finally: with suppress(LookupError): del conn.memory["sherlock"]["futures"]["who_0"][0] users = [] for line in lines: chan, ident, host, server, nick, status, realname = line num_hops, _, realname = realname.partition(' ') users.append((nick, host)) futs = [ wrap_future( event.async_call(update_user_data, db, masks_table, 'mask', now, nick, mask)) for nick, mask in users ] for nick, mask in users: yield from asyncio.gather( ignore_timeout( set_user_data(event, db, hosts_table, 'host', now, nick, get_user_host, conn)), ignore_timeout( set_user_data(event, db, address_table, 'addr', now, nick, get_user_ip, conn))) yield from asyncio.gather(*futs) return "Done."
async def test_send_sieve_error(self, caplog): caplog.set_level(0) conn = make_mock_conn() proto = irc._IrcProtocol(conn) proto._connected = True proto._transport = MagicMock() sieve = object() proto.bot.plugin_manager.out_sieves = [sieve] proto.bot.plugin_manager.internal_launch = launch = MagicMock() fut = async_util.create_future(proto.loop) fut.set_result((False, None)) launch.return_value = fut await proto.send("PRIVMSG #foo bar") assert len(launch.mock_calls) == 1 assert launch.mock_calls[0][1][0] is sieve assert filter_logs(caplog) == [ ( "cloudbot", 30, "Error occurred in outgoing sieve, falling back to old behavior", ), ("cloudbot", 10, "Line was: PRIVMSG #foo bar"), ("cloudbot", 10, "[testconn|out] >> b'PRIVMSG #foo bar\\r\\n'"), ]
def launch(self, hook, event): """ Dispatch a given event to a given hook using a given bot object. Returns False if the hook didn't run successfully, and True if it ran successfully. :type event: cloudbot.event.Event | cloudbot.event.CommandEvent :type hook: cloudbot.plugin.Hook | cloudbot.plugin.CommandHook :rtype: bool """ if hook.type not in ( "on_start", "on_stop", "periodic"): # we don't need sieves on on_start hooks. for sieve in self.bot.plugin_manager.sieves: event = yield from self._sieve(sieve, event, hook) if event is None: return False hook = event.hook if hook.single_thread: # There should only be one running instance of this hook, so let's wait for the last event to be processed # before starting this one. key = (hook.plugin.title, hook.function_name) if key in self._hook_waiting_queues: queue = self._hook_waiting_queues[key] if queue is None: # there's a hook running, but the queue hasn't been created yet, since there's only one hook queue = asyncio.Queue() self._hook_waiting_queues[key] = queue assert isinstance(queue, asyncio.Queue) # create a future to represent this task future = async_util.create_future(self.bot.loop) queue.put_nowait(future) # wait until the last task is completed yield from future else: # set to None to signify that this hook is running, but there's no need to create a full queue # in case there are no more hooks that will wait self._hook_waiting_queues[key] = None # Run the plugin with the message, and wait for it to finish result = yield from self._execute_hook(hook, event) queue = self._hook_waiting_queues[key] if queue is None or queue.empty(): # We're the last task in the queue, we can delete it now. del self._hook_waiting_queues[key] else: # set the result for the next task's future, so they can execute next_future = yield from queue.get() next_future.set_result(None) else: # Run the plugin with the message, and wait for it to finish result = yield from self._execute_hook(hook, event) # Return the result return result
async def on_user_quit(db, event, match): nick = match.group("nick") host = match.group("host") addr = match.group("addr") now = datetime.datetime.now() nick_cf = rfc_casefold(nick) conn = event.conn futures = conn.memory["sherlock"]["futures"] data_futs = futures["data"] nick_change_futs = futures["nick_changes"] try: old_futs = data_futs.pop(nick_cf) except LookupError: nick_change_futs[nick_cf] = fut = create_future(conn.loop) try: old_futs = await asyncio.wait_for(fut, 30) except asyncio.TimeoutError: old_futs = {} finally: with suppress(LookupError): del nick_change_futs[nick_cf] with suppress(KeyError): _set_result(old_futs["host"], host) with suppress(KeyError): _set_result(old_futs["addr"], addr) async def _do_whowas(): mask, _ = await get_user_whowas(event.conn, nick) if mask: with suppress(KeyError): _set_result(old_futs["mask"], mask) await event.async_call( update_user_data, db, masks_table, "mask", now, nick, mask ) await asyncio.gather( event.async_call( update_user_data, db, hosts_table, "host", now, nick, host ), event.async_call( update_user_data, db, address_table, "addr", now, nick, addr ), _do_whowas(), )
def on_user_quit(db, event, match): nick = match.group('nick') ident = match.group('ident') host = match.group('host') addr = match.group('addr') now = datetime.datetime.now() nick_cf = rfc_casefold(nick) conn = event.conn futures = conn.memory["sherlock"]["futures"] data_futs = futures["data"] nick_change_futs = futures["nick_changes"] try: old_futs = data_futs.pop(nick_cf) except LookupError: nick_change_futs[nick_cf] = fut = create_future(conn.loop) try: old_futs = yield from asyncio.wait_for(fut, 30) except asyncio.TimeoutError: old_futs = {} finally: with suppress(LookupError): del nick_change_futs[nick_cf] with suppress(KeyError): _set_result(old_futs['host'], host) with suppress(KeyError): _set_result(old_futs['addr'], addr) @asyncio.coroutine def _do_whowas(): mask, _ = yield from get_user_whowas(event.conn, nick) if mask: with suppress(KeyError): _set_result(old_futs['mask'], mask) yield from event.async_call(update_user_data, db, masks_table, 'mask', now, nick, mask) yield from asyncio.gather( event.async_call(update_user_data, db, hosts_table, 'host', now, nick, host), event.async_call(update_user_data, db, address_table, 'addr', now, nick, addr), _do_whowas(), )
async def test_send_sieve_error(self): proto = irc._IrcProtocol(MagicMock(loop=asyncio.get_event_loop())) proto._connected = True proto._transport = MagicMock() sieve = object() proto.bot.plugin_manager.out_sieves = [sieve] proto.bot.plugin_manager.internal_launch = launch = MagicMock() fut = async_util.create_future(proto.loop) fut.set_result((False, None)) launch.return_value = fut await proto.send("PRIVMSG #foo bar") assert len(launch.mock_calls) == 1 assert launch.mock_calls[0][1][0] is sieve
def get_user_whowas(conn, nick): nick_cf = rfc_casefold(nick) futs = conn.memory["sherlock"]["futures"] send_command = False try: mask_fut = futs["user_whowas_mask"][nick_cf] except LookupError: send_command = True futs["user_whowas_mask"][nick_cf] = mask_fut = create_future(conn.loop) try: host_fut = futs["user_whowas_host"][nick_cf] except LookupError: send_command = True futs["user_whowas_host"][nick_cf] = host_fut = create_future(conn.loop) if send_command: conn.cmd("WHOWAS", nick) try: return (yield from await_response(asyncio.gather(mask_fut, host_fut))) except asyncio.TimeoutError: return None, None
async def get_user_whowas(conn, nick): nick_cf = rfc_casefold(nick) futs = conn.memory["sherlock"]["futures"] send_command = False try: mask_fut = futs["user_whowas_mask"][nick_cf] except LookupError: send_command = True futs["user_whowas_mask"][nick_cf] = mask_fut = create_future(conn.loop) try: host_fut = futs["user_whowas_host"][nick_cf] except LookupError: send_command = True futs["user_whowas_host"][nick_cf] = host_fut = create_future(conn.loop) if send_command: conn.cmd("WHOWAS", nick) try: return tuple(await await_response(asyncio.gather(mask_fut, host_fut))) except asyncio.TimeoutError: logger.warning("[user_tracking] Hit timeout for whowas data") return None, None
async def await_command_response(conn, name, cmd, nick): futs = conn.memory["sherlock"]["futures"][name] nick_cf = rfc_casefold(nick) try: fut = futs[nick_cf] except LookupError: futs[nick_cf] = fut = create_future(conn.loop) conn.cmd(cmd, nick) try: res = await await_response(fut) finally: with suppress(KeyError): del futs[nick_cf] return res
def on_user_connect(db, event, match): nick = match.group('nick') ident = match.group('ident') host = match.group('host') addr = match.group('addr') now = datetime.datetime.now() nick_cf = rfc_casefold(nick) conn = event.conn mask_fut = create_future(conn.loop) futs = { 'mask': mask_fut, } data_futs = conn.memory["sherlock"]["futures"]["data"] data_futs[nick_cf] = futs @asyncio.coroutine def _handle_set(table, name, value_func): try: value = yield from value_func(event.conn, nick) except (asyncio.TimeoutError, asyncio.CancelledError): value = yield from futs[name] del futs[name] yield from event.async_call(update_user_data, db, table, name, now, nick, value) @asyncio.coroutine def _do_mask(): yield from _handle_set(masks_table, 'mask', get_user_mask) yield from asyncio.gather( event.async_call(update_user_data, db, hosts_table, 'host', now, nick, host), event.async_call(update_user_data, db, address_table, 'addr', now, nick, addr), _do_mask(), ) with suppress(LookupError): del data_futs[nick_cf]
async def on_user_connect(db, event, match): nick = match.group("nick") host = match.group("host") addr = match.group("addr") now = datetime.datetime.now() nick_cf = rfc_casefold(nick) conn = event.conn mask_fut = create_future(conn.loop) futs = { "mask": mask_fut, } data_futs = conn.memory["sherlock"]["futures"]["data"] data_futs[nick_cf] = futs async def _handle_set(table, name, value_func): try: value = await value_func(event.conn, nick) except (asyncio.TimeoutError, asyncio.CancelledError): value = await futs[name] del futs[name] await event.async_call( update_user_data, db, table, name, now, nick, value ) async def _do_mask(): await _handle_set(masks_table, "mask", get_user_mask) await asyncio.gather( event.async_call( update_user_data, db, hosts_table, "host", now, nick, host ), event.async_call( update_user_data, db, address_table, "addr", now, nick, addr ), _do_mask(), ) with suppress(LookupError): del data_futs[nick_cf]
def __init__(self, max_executors=None, executor_type=ThreadPoolExecutor, **kwargs): if max_executors is None: max_executors = (os.cpu_count() or 1) * 5 if max_executors <= 0: raise ValueError("max_executors must be greater than 0") self._max = max_executors self._exec_class = executor_type self._exec_args = kwargs self._executors = [] self._free_executors = [] self._executor_waiter = create_future()
def __init__(self, conn): """ :type conn: IrcClient """ self.loop = conn.loop self.bot = conn.bot self.conn = conn # input buffer self._input_buffer = b"" # connected self._connected = False self._connecting = True # transport self._transport = None # Future that waits until we are connected self._connected_future = async_util.create_future(self.loop)
def __init__(self, bot, _type, name, nick, *, channels=None, config=None): """ :type bot: cloudbot.bot.CloudBot :type name: str :type nick: str :type channels: list[str] :type config: dict[str, unknown] """ self.bot = bot self.loop = bot.loop self.name = name self.nick = nick self._type = _type self.channels = [] if channels is None: self.config_channels = [] else: self.config_channels = channels if config is None: self.config = {} else: self.config = config self.vars = {} self.history = {} # create permissions manager self.permissions = PermissionManager(self) # for plugins to abuse self.memory = collections.defaultdict() # set when on_load in core_misc is done self.ready = False self._active = False self.cancelled_future = async_util.create_future(self.loop)
async def handle_available_caps(conn, caplist, event, irc_paramlist, bot): available_caps = conn.memory["available_caps"] # type: CapList available_caps.extend(caplist) cap_queue = conn.memory["cap_queue"] for cap in caplist: name = cap.name name_cf = name.casefold() cap_event = partial(CapEvent, base_event=event, cap=name, cap_param=cap.value) tasks = [ bot.plugin_manager.internal_launch(_hook, cap_event(hook=_hook)) for _hook in bot.plugin_manager.cap_hooks["on_available"][name_cf] ] results = await asyncio.gather(*tasks) if any(ok and (res or res is None) for ok, res in results): cap_queue[name_cf] = async_util.create_future(conn.loop) conn.cmd("CAP", "REQ", name) if irc_paramlist[2] != '*': await asyncio.gather(*cap_queue.values()) cap_queue.clear() conn.send("CAP END")
def __init__(self, loop=asyncio.get_event_loop()): if bot.get(): raise ValueError("There seems to already be a bot running!") bot.set(self) # basic variables self.base_dir = Path().resolve() self.loop = loop self.start_time = time.time() self.running = True self.clients = {} # future which will be called when the bot stopsIf you self.stopped_future = async_util.create_future(self.loop) # stores each bot server connection self.connections = KeyFoldDict() # for plugins self.logger = logger # for plugins to abuse self.memory = collections.defaultdict() # declare and create data folder self.data_dir = os.path.abspath('data') if not os.path.exists(self.data_dir): logger.debug("Data folder not found, creating.") os.mkdir(self.data_dir) # set up config self.config = Config(self) logger.debug("Config system initialised.") # set values for reloading reloading_conf = self.config.get("reloading", {}) self.plugin_reloading_enabled = reloading_conf.get( "plugin_reloading", False) self.config_reloading_enabled = reloading_conf.get( "config_reloading", True) # this doesn't REALLY need to be here but it's nice self.user_agent = self.config.get( 'user_agent', 'CloudBot/3.0 - CloudBot Refresh ' '<https://github.com/CloudBotIRC/CloudBot/>') # setup db db_path = self.config.get('database', 'sqlite:///cloudbot.db') self.db_engine = create_engine(db_path) self.db_factory = sessionmaker(bind=self.db_engine) self.db_session = scoped_session(self.db_factory) self.db_metadata = database.metadata self.db_base = declarative_base(metadata=self.db_metadata, bind=self.db_engine) self.db_executor_pool = ExecutorPool(50, max_workers=1, thread_name_prefix='cloudbot-db') # set botvars so plugins can access when loading database.base = self.db_base logger.debug("Database system initialised.") # Bot initialisation complete logger.debug("Bot setup completed.") self.load_clients() # create bot connections self.create_connections() self.observer = Observer() if self.plugin_reloading_enabled: self.plugin_reloader = PluginReloader(self) if self.config_reloading_enabled: self.config_reloader = ConfigReloader(self) self.plugin_manager = PluginManager(self)
async def on_nickchange(db, event, match): conn = event.conn old_nick = match.group("oldnick") new_nick = match.group("newnick") now = datetime.datetime.now() old_nick_cf = rfc_casefold(old_nick) new_nick_cf = rfc_casefold(new_nick) futures = conn.memory["sherlock"]["futures"] data_futs = futures["data"] nick_change_futs = futures["nick_changes"] futs = { "addr": create_future(conn.loop), "host": create_future(conn.loop), "mask": create_future(conn.loop), } data_futs[new_nick_cf] = futs try: old_futs = data_futs.pop(old_nick_cf) except LookupError: nick_change_futs[old_nick_cf] = fut = create_future(conn.loop) try: old_futs = await asyncio.wait_for(fut, 30) except asyncio.TimeoutError: old_futs = {} finally: with suppress(LookupError): del nick_change_futs[old_nick_cf] try: nick_fut = nick_change_futs.pop(new_nick_cf) except LookupError: pass else: _set_result(nick_fut, futs) async def _handle_set(table, name, value_func): try: value = await value_func(event.conn, new_nick) except (asyncio.TimeoutError, asyncio.CancelledError): value = await asyncio.wait_for(futs[name], 300) del futs[name] with suppress(LookupError): _set_result(old_futs[name], value) await asyncio.gather( event.async_call( update_user_data, db, table, name, now, old_nick, value ), event.async_call( update_user_data, db, table, name, now, new_nick, value ), ) async def _do_mask(): await _handle_set(masks_table, "mask", get_user_mask) async def _timeout_whowas(): try: await asyncio.gather( _handle_set(hosts_table, "host", get_user_host), _do_mask(), ) except (asyncio.TimeoutError, asyncio.CancelledError): mask, host = await get_user_whowas(event.conn, new_nick) if mask and host: with suppress(KeyError): _set_result(old_futs["host"], host) with suppress(KeyError): _set_result(old_futs["mask"], mask) await asyncio.gather( event.async_call( update_user_data, db, hosts_table, "host", now, old_nick, host, ), event.async_call( update_user_data, db, hosts_table, "host", now, new_nick, host, ), event.async_call( update_user_data, db, masks_table, "mask", now, old_nick, mask, ), event.async_call( update_user_data, db, masks_table, "mask", now, new_nick, mask, ), ) await asyncio.gather( _handle_set(address_table, "addr", get_user_ip), _timeout_whowas(), ) with suppress(KeyError): del data_futs[new_nick_cf]
def on_nickchange(db, event, match): conn = event.conn old_nick = match.group('oldnick') new_nick = match.group('newnick') now = datetime.datetime.now() old_nick_cf = rfc_casefold(old_nick) new_nick_cf = rfc_casefold(new_nick) futures = conn.memory["sherlock"]["futures"] data_futs = futures["data"] nick_change_futs = futures["nick_changes"] futs = { 'addr': create_future(conn.loop), 'host': create_future(conn.loop), 'mask': create_future(conn.loop), } data_futs[new_nick_cf] = futs try: old_futs = data_futs.pop(old_nick_cf) except LookupError: nick_change_futs[old_nick_cf] = fut = create_future(conn.loop) try: old_futs = yield from asyncio.wait_for(fut, 30) except asyncio.TimeoutError: old_futs = {} finally: with suppress(LookupError): del nick_change_futs[old_nick_cf] try: nick_fut = nick_change_futs.pop(new_nick_cf) except LookupError: pass else: _set_result(nick_fut, futs) @asyncio.coroutine def _handle_set(table, name, value_func): try: value = yield from value_func(event.conn, new_nick) except (asyncio.TimeoutError, asyncio.CancelledError): value = yield from asyncio.wait_for(futs[name], 300) del futs[name] with suppress(LookupError): _set_result(old_futs[name], value) yield from asyncio.gather( event.async_call(update_user_data, db, table, name, now, old_nick, value), event.async_call(update_user_data, db, table, name, now, new_nick, value)) @asyncio.coroutine def _do_mask(): yield from _handle_set(masks_table, 'mask', get_user_mask) @asyncio.coroutine def _timeout_whowas(): try: yield from asyncio.gather( _handle_set(hosts_table, 'host', get_user_host), _do_mask(), ) except (asyncio.TimeoutError, asyncio.CancelledError): mask, host = yield from get_user_whowas(event.conn, new_nick) if mask and host: with suppress(KeyError): _set_result(old_futs['host'], host) with suppress(KeyError): _set_result(old_futs['mask'], mask) yield from asyncio.gather( event.async_call(update_user_data, db, hosts_table, 'host', now, old_nick, host), event.async_call(update_user_data, db, hosts_table, 'host', now, new_nick, host), event.async_call(update_user_data, db, masks_table, 'mask', now, old_nick, mask), event.async_call(update_user_data, db, masks_table, 'mask', now, new_nick, mask), ) yield from asyncio.gather( _handle_set(address_table, 'addr', get_user_ip), _timeout_whowas(), ) with suppress(KeyError): del data_futs[new_nick_cf]