""" .. module:: IRCCore :synopsis: IRC protocol class .. moduleauthor:: Dave Pedu """ import socket import asyncio import logging import traceback import sys from inspect import getargspec from pyircbot.common import burstbucket, parse_irc_line from collections import namedtuple from io import StringIO IRCEvent = namedtuple("IRCEvent", "command args prefix trailing") UserPrefix = namedtuple("UserPrefix", "nick username hostname") ServerPrefix = namedtuple("ServerPrefix", "hostname") class IRCCore(object): def __init__(self, servers, loop, rate_limit=True, rate_max=5.0, rate_int=1.1): self._loop = loop # rate limiting options self.rate_limit = rate_limit self.rate_max = float(rate_max) self.rate_int = float(rate_int) self.reconnect_delay = 3.0 self.connected = False """If we're connected or not""" self.log = logging.getLogger('IRCCore') """Reference to logger object""" self.buffer = StringIO() """cStringIO used as a buffer""" self.alive = True """True if we should try to stay connected""" self.server = 0 """Current server index""" self.servers = servers """List of server address""" self.port = 0 """Server port""" self.connection_family = socket.AF_UNSPEC """Socket family. 0 will auto-detect ipv4 or v6. Change this to socket.AF_INET or socket.AF_INET6 force use of ipv4 or ipv6.""" self.bind_addr = None """Optionally bind to a specific address. This should be a (host, port) tuple.""" self.nick = None # Set up hooks for modules self.initHooks() self.outseq = 5 self.outputq = asyncio.PriorityQueue() self._loop.call_soon_threadsafe(asyncio.ensure_future, self.outputqueue()) async def loop(self, loop): while self.alive: try: # TODO support ipv6 again self.reader, self.writer = await asyncio.open_connection(self.servers[self.server][0], port=self.servers[self.server][1], loop=loop, ssl=None, family=self.connection_family, local_addr=self.bind_addr) self.fire_hook("_CONNECT") except (socket.gaierror, ConnectionRefusedError): traceback.print_exc() logging.warning("Non-fatal connect error, trying next server...") self.server = (self.server + 1) % len(self.servers) await asyncio.sleep(1, loop=loop) continue while self.alive: try: data = await self.reader.readuntil() self.log.debug("<<< {}".format(repr(data))) command, args, prefix, trailing = parse_irc_line(data.decode("UTF-8")) self.fire_hook("_RECV", args=args, prefix=prefix, trailing=trailing) if command not in self.hookcalls: self.log.warning("Unknown command: cmd='{}' prefix='{}' args='{}' trailing='{}'" .format(command, prefix, args, trailing)) else: self.fire_hook(command, args=args, prefix=prefix, trailing=trailing) except (ConnectionResetError, asyncio.streams.IncompleteReadError): traceback.print_exc() break except (UnicodeDecodeError, ): traceback.print_exc() self.fire_hook("_DISCONNECT") self.writer.close() if self.alive: # TODO ramp down reconnect attempts logging.info("Reconnecting in {}s...".format(self.reconnect_delay)) await asyncio.sleep(self.reconnect_delay) async def outputqueue(self): self.bucket = burstbucket(self.rate_max, self.rate_int) while True: prio, line = await self.outputq.get() # sleep until the bucket allows us to send if self.rate_limit: while True: s = self.bucket.get() if s == 0: break else: await asyncio.sleep(s, loop=self._loop) self.fire_hook('_SEND', args=None, prefix=None, trailing=None) self.log.debug(">>> {}".format(repr(line))) self.outputq.task_done() try: self.writer.write((line + "\r\n").encode("UTF-8")) except Exception as e: # Probably fine if we drop messages while offline print(e) print(self.trace()) async def kill(self, message="Help! Another thread is killing me :(", forever=True): """Send quit message, flush queue, and close the socket :param message: Quit message to send before disconnecting :type message: str """ if forever: self.alive = False self.act_QUIT(message) # TODO will this hang if the socket is having issues? await self.writer.drain() self.writer.close() self.log.info("Kill complete") def sendRaw(self, data, priority=None): """ Send data on the wire. Lower priorities are sent first. :param data: unicode data to send. will be converted to utf-8 :param priority: numerical priority value. If not None, the message will likely be sent first. Otherwise, an ever-increasing sequence number is used to maintain order. For a minimum priority message, use a priority value of sys.maxsize. """ if priority is None: self.outseq += 1 priority = self.outseq asyncio.run_coroutine_threadsafe(self.outputq.put((priority, data, )), self._loop) " Module related code " def initHooks(self): """Defines hooks that modules can listen for events of""" self.hooks = [ '_ALL', '_CONNECT', '_DISCONNECT', '_RECV', '_SEND', 'NOTICE', 'MODE', 'PING', 'JOIN', 'QUIT', 'NICK', 'PART', 'PRIVMSG', 'KICK', 'INVITE', '001', '002', '003', '004', '005', '250', '251', '252', '254', '255', '265', '266', '331', '332', '333', '353', '366', '372', '375', '376', '401', '422', '433', ] " mapping of hooks to methods " self.hookcalls = {command: [] for command in self.hooks} def fire_hook(self, command, args=None, prefix=None, trailing=None): """Run any listeners for a specific hook :param command: the hook to fire :type command: str :param args: the list of arguments, if any, the command was passed :type args: list :param prefix: prefix of the sender of this command :type prefix: str :param trailing: data payload of the command :type trailing: str""" for hook in self.hookcalls["_ALL"] + self.hookcalls[command]: try: if len(getargspec(hook).args) == 2: hook(IRCCore.packetAsObject(command, args, prefix, trailing)) else: hook(args, prefix, trailing) except: self.log.warning("Error processing hook: \n%s" % self.trace()) def addHook(self, command, method): """**Internal.** Enable (connect) a single hook of a module :param command: command this hook will trigger on :type command: str :param method: callable method object to hook in :type method: object""" " add a single hook " if command in self.hooks: self.hookcalls[command].append(method) else: self.log.warning("Invalid hook - %s" % command) return False def removeHook(self, command, method): """**Internal.** Disable (disconnect) a single hook of a module :param command: command this hook triggers on :type command: str :param method: callable method that should be removed :type method: object""" " remove a single hook " if command in self.hooks: for hookedMethod in self.hookcalls[command]: if hookedMethod == method: self.hookcalls[command].remove(hookedMethod) else: self.log.warning("Invalid hook - %s" % command) return False def packetAsObject(command, args, prefix, trailing): """Given an irc message's args, prefix, and trailing data return an object with these properties :param args: list of args from the IRC packet :type args: list :param prefix: prefix object parsed from the IRC packet :type prefix: ServerPrefix or UserPrefix :param trailing: trailing data from the IRC packet :type trailing: str :returns: object -- a IRCEvent object with the ``args``, ``prefix``, ``trailing``""" return IRCEvent(command, args, IRCCore.decodePrefix(prefix) if prefix else None, trailing) " Utility methods " @staticmethod def decodePrefix(prefix): """Given a prefix like nick!username@hostname, return an object with these properties :param prefix: the prefix to disassemble :type prefix: str :returns: object -- an UserPrefix object with the properties `nick`, `username`, `hostname` or a ServerPrefix object with the property `hostname` """ if "!" in prefix: nick, prefix = prefix.split("!") username, hostname = prefix.split("@") return UserPrefix(nick, username, hostname) else: return ServerPrefix(prefix) @staticmethod def trace(): """Return the stack trace of the bot as a string""" return traceback.format_exc() @staticmethod def fulltrace(): """Return the stack trace of the bot as a string""" result = "" result += "\n*** STACKTRACE - START ***\n" code = [] for threadId, stack in sys._current_frames().items(): code.append("\n# ThreadID: %s" % threadId) for filename, lineno, name, line in traceback.extract_stack(stack): code.append('File: "%s", line %d, in %s' % (filename, lineno, name)) if line: code.append(" %s" % (line.strip())) for line in code: result += line + "\n" result += "\n*** STACKTRACE - END ***\n" return result " Data Methods " def get_nick(self): """Get the bot's current nick :returns: str - the bot's current nickname""" return self.nick " Action Methods " def act_PONG(self, data): """Use the `/pong` command - respond to server pings :param data: the string or number the server sent with it's ping :type data: str""" self.sendRaw("PONG :%s" % data) def act_USER(self, username, hostname, realname, priority=2): """Use the USER protocol command. Used during connection :param username: the bot's username :type username: str :param hostname: the bot's hostname :type hostname: str :param realname: the bot's realname :type realname: str""" self.sendRaw("USER %s %s %s :%s" % (username, hostname, self.servers[self.server], realname), priority) def act_NICK(self, newNick, priority=2): """Use the `/nick` command :param newNick: new nick for the bot :type newNick: str""" self.nick = newNick self.sendRaw("NICK %s" % newNick, priority) def act_JOIN(self, channel, priority=3): """Use the `/join` command :param channel: the channel to attempt to join :type channel: str""" self.sendRaw("JOIN %s" % channel, priority=3) def act_PRIVMSG(self, towho, message): """Use the `/msg` command :param towho: the target #channel or user's name :type towho: str :param message: the message to send :type message: str""" self.sendRaw("PRIVMSG %s :%s" % (towho, message)) def act_MODE(self, channel, mode, extra=None): """Use the `/mode` command :param channel: the channel this mode is for :type channel: str :param mode: the mode string. Example: +b :type mode: str :param extra: additional argument if the mode needs it. Example: user@*!* :type extra: str""" if extra is not None: self.sendRaw("MODE %s %s %s" % (channel, mode, extra)) else: self.sendRaw("MODE %s %s" % (channel, mode)) def act_ACTION(self, channel, action): """Use the `/me ` command :param channel: the channel name or target's name the message is sent to :type channel: str :param action: the text to send :type action: str""" self.sendRaw("PRIVMSG %s :\x01ACTION %s" % (channel, action)) def act_KICK(self, channel, who, comment=""): """Use the `/kick ` command :param channel: the channel from which the user will be kicked :type channel: str :param who: the nickname of the user to kick :type action: str :param comment: the kick message :type comment: str""" self.sendRaw("KICK %s %s :%s" % (channel, who, comment)) def act_QUIT(self, message, priority=2): """Use the `/quit` command :param message: quit message :type message: str""" self.sendRaw("QUIT :%s" % message, priority) def act_PASS(self, password): """ Send server password, for use on connection """ self.sendRaw("PASS %s" % password)