Awesome IRC bot
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.

403 lines
14 KiB

"""
.. module:: IRCCore
:synopsis: IRC protocol class
.. moduleauthor:: Dave Pedu <dave@davepedu.com>
"""
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 <action>` 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 <user> <message>` 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)