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.

534 lines
20 KiB

"""
.. module:: IRCCore
:synopsis: IRC protocol class
.. moduleauthor:: Dave Pedu <dave@davepedu.com>
"""
7 years ago
import queue
import socket
import asynchat
import asyncore
import logging
import traceback
import sys
from inspect import getargspec
from socket import SHUT_RDWR
7 years ago
from threading import Thread
from time import sleep, time
from collections import namedtuple
try:
from cStringIO import StringIO
except:
from io import BytesIO as StringIO
IRCEvent = namedtuple("IRCEvent", "args prefix trailing")
UserPrefix = namedtuple("UserPrefix", "nick username hostname")
ServerPrefix = namedtuple("ServerPrefix", "hostname")
class IRCCore(asynchat.async_chat):
def __init__(self):
asynchat.async_chat.__init__(self)
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 = []
"""List of server address"""
self.port = 0
"""Server port"""
self.ipv6 = False
"""Use IPv6?"""
7 years ago
self.OUTPUT_BUFFER_SIZE = 1000
self.SEND_WAIT = 0.800
6 years ago
self.outputQueue = queue.Queue(self.OUTPUT_BUFFER_SIZE)
7 years ago
self.outputQueueRunner = OutputQueueRunner(self)
self.outputQueueRunner.start()
# IRC Messages are terminated with \r\n
self.set_terminator(b"\r\n")
# Set up hooks for modules
self.initHooks()
# Map for asynchat
self.asynmap = {}
def loop(self):
while self.alive:
try:
asyncore.loop(map=self.asynmap, timeout=1)
except Exception:
self.log.error("Loop error: ")
self.log.error(IRCCore.trace())
# Remove from asynmap
for key in list(self.asynmap.keys())[:]:
del self.asynmap[key]
if self.alive:
logging.info("Loop: reconnecting")
try:
self._connect()
except Exception:
self.log.error("Error reconnecting: ")
self.log.error(IRCCore.trace())
def kill(self, message="Help! Another thread is killing me :(", alive=False):
"""Send quit message, flush queue, and close the socket
:param message: Quit message
:type message: str
:param alive: True causes a reconnect after disconnecting
:type alive: bool
"""
7 years ago
# Pauses output queue
self.outputQueueRunner.paused = not alive
7 years ago
# Clear any pending messages
self.outputQueueRunner.clear()
# Send quit message and flush queue
self.act_QUIT(message) # TODO will this hang if the socket is having issues?
7 years ago
self.outputQueueRunner.flush()
# Signal disconnection
self.alive = alive
7 years ago
# Close socket
7 years ago
self.socket.shutdown(SHUT_RDWR)
7 years ago
self.close()
self.log.info("Kill complete")
" Net related code here on down "
def getBuf(self):
"""Return the network buffer and clear it"""
self.buffer.seek(0)
data = self.buffer.read()
self.buffer = StringIO()
return data
def collect_incoming_data(self, data):
"""Recieve data from the IRC server, append it to the buffer
:param data: the data that was recieved
:type data: str"""
self.buffer.write(data)
def found_terminator(self):
"""A complete command was pushed through, so clear the buffer and process it."""
line = None
buf = self.getBuf()
try:
line = buf.decode("UTF-8")
except UnicodeDecodeError as ude:
self.log.error("found_terminator(): could not decode input as UTF-8")
self.log.error("found_terminator(): data: %s" % line)
self.log.error("found_terminator(): repr(data): %s" % repr(line))
self.log.error("found_terminator(): error: %s" % str(ude))
return
self.log.debug("< {}".format(line))
self.process_data(line)
def handle_close(self):
"""Called when the socket is disconnected. Triggers the _DISCONNECT hook"""
self.log.info("handle_close")
self.connected = False
self.close()
self.fire_hook("_DISCONNECT")
def handle_error(self, *args, **kwargs):
"""Called on fatal network errors."""
self.log.error("Connection failed (handle_error)")
self.log.error(str(args))
self.log.error(str(kwargs))
self.log.error(IRCCore.trace())
def _connect(self):
"""Connect to IRC"""
self.server += 1
if self.server >= len(self.servers):
self.server = 0
serverHostname = self.servers[self.server]
self.log.info("Connecting to %(server)s:%(port)i", {"server": serverHostname, "port": self.port})
socket_type = socket.AF_INET
if self.ipv6:
self.log.info("IPv6 is enabled.")
socket_type = socket.AF_INET6
socketInfo = socket.getaddrinfo(serverHostname, self.port, socket_type)
self.create_socket(socket_type, socket.SOCK_STREAM)
self.log.info("Socket created: %s" % self.socket.fileno())
self.connect(socketInfo[0][4])
self.log.info("Connection established")
self._fileno = self.socket.fileno()
# See http://willpython.blogspot.com/2010/08/multiple-event-loops-with-asyncore-and.html
self.asynmap[self._fileno] = self
self.log.info("_connect: Socket map: %s" % str(self.asynmap))
def handle_connect(self):
"""When asynchat indicates our socket is connected, fire the _CONNECT hook"""
self.connected = True
self.log.info("handle_connect: connected")
self.fire_hook("_CONNECT")
self.log.info("handle_connect: complete")
7 years ago
def sendRaw(self, text, prio=2):
"""Queue messages (raw string) to be sent to the IRC server
:param text: the string to send
:type text: str"""
text = (text + "\r\n").encode("UTF-8").decode().encode("UTF-8")
7 years ago
self.outputQueue.put((prio, text), block=False)
def process_data(self, data):
"""Process one line of tet irc sent us
:param data: the data to process
:type data: str"""
if data.strip() == "":
return
prefix = None
command = None
args = []
trailing = None
if data[0] == ":":
prefix = data.split(" ")[0][1:]
data = data[data.find(" ") + 1:]
command = data.split(" ")[0]
data = data[data.find(" ") + 1:]
if(data[0] == ":"):
# no args
trailing = data[1:].strip()
else:
trailing = data[data.find(" :") + 2:].strip()
data = data[:data.find(" :")]
args = data.split(" ")
for index, arg in enumerate(args):
args[index] = arg.strip()
7 years ago
self.fire_hook("_RECV", args=args, prefix=prefix, trailing=trailing)
if command not in self.hookcalls:
self.log.warning("Unknown command: cmd='%s' prefix='%s' args='%s' trailing='%s'" % (command, prefix, args,
trailing))
else:
self.fire_hook(command, args=args, prefix=prefix, trailing=trailing)
" Module related code "
def initHooks(self):
"""Defines hooks that modules can listen for events of"""
self.hooks = [
'_CONNECT', # Called when the bot connects to IRC on the socket level
'_DISCONNECT', # Called when the irc socket is forcibly closed
'_RECV', # Called on network activity
'NOTICE', # :irc.129irc.com NOTICE AUTH :*** Looking up your hostname...
'MODE', # :CloneABCD MODE CloneABCD :+iwx
'PING', # PING :irc.129irc.com
'JOIN', # :CloneA!dave@hidden-B4F6B1AA.rit.edu JOIN :#clonea
'QUIT', # :HCSMPBot!~HCSMPBot@108.170.48.18 QUIT :Quit: Disconnecting!
'NICK', # :foxiAway!foxi@irc.hcsmp.com NICK :foxi
'PART', # :CloneA!dave@hidden-B4F6B1AA.rit.edu PART #clonea
'PRIVMSG', # :CloneA!dave@hidden-B4F6B1AA.rit.edu PRIVMSG #clonea :aaa
'KICK', # :xMopxShell!~rduser@host KICK #xMopx2 xBotxShellTest :xBotxShellTest
'INVITE', # :gmx!~gmxgeek@irc.hcsmp.com INVITE Tyrone :#hcsmp'
'001', # :irc.129irc.com 001 CloneABCD :Welcome to the 129irc IRC Network CloneABCD!CloneABCD@djptwc-laptop1.rit.edu
'002', # :irc.129irc.com 002 CloneABCD :Your host is irc.129irc.com, running version Unreal3.2.8.1
'003', # :irc.129irc.com 003 CloneABCD :This server was created Mon Jul 19 2010 at 03:12:01 EDT
'004', # :irc.129irc.com 004 CloneABCD irc.129irc.com Unreal3.2.8.1 iowghraAsORTVSxNCWqBzvdHtGp lvhopsmntikrRcaqOALQbSeIKVfMCuzNTGj
'005', # :irc.129irc.com 005 CloneABCD CMDS=KNOCK,MAP,DCCALLOW,USERIP UHNAMES NAMESX SAFELIST HCN MAXCHANNELS=10 CHANLIMIT=#:10 MAXLIST=b:60,e:60,I:60 NICKLEN=30 CHANNELLEN=32 TOPICLEN=307 KICKLEN=307 AWAYLEN=307 :are supported by this server
'250', # :chaos.esper.net 250 xBotxShellTest :Highest connection count: 1633 (1632 clients) (186588 connections received)
'251', # :irc.129irc.com 251 CloneABCD :There are 1 users and 48 invisible on 2 servers
'252', # :irc.129irc.com 252 CloneABCD 9 :operator(s) online
'254', # :irc.129irc.com 254 CloneABCD 6 :channels formed
'255', # :irc.129irc.com 255 CloneABCD :I have 42 clients and 1 servers
'265', # :irc.129irc.com 265 CloneABCD :Current Local Users: 42 Max: 47
'266', # :irc.129irc.com 266 CloneABCD :Current Global Users: 49 Max: 53
'332', # :chaos.esper.net 332 xBotxShellTest #xMopx2 :/ #XMOPX2 / https://code.google.com/p/pyircbot/ (Channel Topic)
'333', # :chaos.esper.net 333 xBotxShellTest #xMopx2 xMopxShell!~rduser@108.170.60.242 1344370109
'353', # :irc.129irc.com 353 CloneABCD = #clonea :CloneABCD CloneABC
'366', # :irc.129irc.com 366 CloneABCD #clonea :End of /NAMES list.
'372', # :chaos.esper.net 372 xBotxShell :motd text here
'375', # :chaos.esper.net 375 xBotxShellTest :- chaos.esper.net Message of the Day -
'376', # :chaos.esper.net 376 xBotxShell :End of /MOTD command.
'422', # :irc.129irc.com 422 CloneABCD :MOTD File is missing
'433', # :nova.esper.net 433 * pyircbot3 :Nickname is already in use.
]
" 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[command]:
try:
if len(getargspec(hook).args) == 2:
hook(IRCCore.packetAsObject(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(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(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):
"""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))
def act_NICK(self, newNick):
"""Use the `/nick` command
:param newNick: new nick for the bot
:type newNick: str"""
self.nick = newNick
self.sendRaw("NICK %s" % newNick)
def act_JOIN(self, channel):
"""Use the `/join` command
:param channel: the channel to attempt to join
:type channel: str"""
self.sendRaw("JOIN %s" % channel)
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):
"""Use the `/quit` command
:param message: quit message
:type message: str"""
7 years ago
self.sendRaw("QUIT :%s" % message, prio=0)
def act_PASS(self, password):
"""
Send server password, for use on connection
"""
self.sendRaw("PASS %s" % password, prio=0)
7 years ago
class OutputQueueRunner(Thread):
"""Rate-limited output queue"""
def __init__(self, bot):
Thread.__init__(self, daemon=True)
self.bot = bot # reference to main bot thread
7 years ago
self.log = logging.getLogger('OutputQueueRunner')
self.paused = False
7 years ago
def run(self):
"""Constantly sends messages unless bot is disconnecting/ed"""
lastSend = time()
while True:
# Rate limit
sinceLast = time() - lastSend
if sinceLast < self.bot.SEND_WAIT:
toSleep = self.bot.SEND_WAIT - sinceLast
sleep(toSleep)
7 years ago
# Pop item and execute
if self.bot.connected and not self.paused:
try:
self.process_queue_item()
lastSend = time()
except queue.Empty:
# self.log.debug("Queue is empty")
7 years ago
pass
sleep(0.01)
7 years ago
def process_queue_item(self):
"""Remove 1 item from queue and process it"""
prio, text = self.bot.outputQueue.get(block=True, timeout=10)
# self.log.debug("%s>> %s" % (prio,text))
7 years ago
self.bot.outputQueue.task_done()
self.log.debug("> {}".format(text.decode('UTF-8')))
7 years ago
self.bot.send(text)
7 years ago
def clear(self):
"""Discard all items from queue"""
length = self.bot.outputQueue.qsize()
try:
while True:
self.bot.outputQueue.get(block=False)
except queue.Empty:
pass
# self.log.debug("output queue cleared")
7 years ago
return length
7 years ago
def flush(self):
"""Process all items in queue"""
for i in range(0, self.bot.outputQueue.qsize()):
try:
self.process_queue_item()
except:
pass
# self.log.debug("output queue flushed")