3
0
mirror of https://github.com/jlu5/PyLink.git synced 2024-11-01 01:09:22 +01:00
PyLink/classes.py
James Lu eac81de176 Irc: add internal support for aborting connections
New "aborted" threading event, which can be set to True in order to stop the existing run() loop

Autoconnect confirmed to work with this.
2015-09-02 22:06:18 -07:00

403 lines
16 KiB
Python

import threading
from random import randint
import time
import socket
import threading
import ssl
from collections import defaultdict
import hashlib
from log import log
from conf import conf
import world
### Exceptions
class ProtocolError(Exception):
pass
### Internal classes (users, servers, channels)
class Irc():
def initVars(self):
self.connected.clear()
self.aborted.clear()
self.pseudoclient = None
self.lastping = time.time()
# Server, channel, and user indexes to be populated by our protocol module
self.servers = {self.sid: IrcServer(None, self.serverdata['hostname'], internal=True)}
self.users = {}
self.channels = defaultdict(IrcChannel)
# Sets flags such as whether to use halfops, etc. The default RFC1459
# modes are implied.
self.cmodes = {'op': 'o', 'secret': 's', 'private': 'p',
'noextmsg': 'n', 'moderated': 'm', 'inviteonly': 'i',
'topiclock': 't', 'limit': 'l', 'ban': 'b',
'voice': 'v', 'key': 'k',
# Type A, B, and C modes
'*A': 'b',
'*B': 'k',
'*C': 'l',
'*D': 'imnpstr'}
self.umodes = {'invisible': 'i', 'snomask': 's', 'wallops': 'w',
'oper': 'o',
'*A': '', '*B': '', '*C': 's', '*D': 'iow'}
# This max nick length starts off as the config value, but may be
# overwritten later by the protocol module if such information is
# received. Note that only some IRCds (InspIRCd) give us nick length
# during link, so it is still required that the config value be set!
self.maxnicklen = self.serverdata['maxnicklen']
self.prefixmodes = {'o': '@', 'v': '+'}
# Uplink SID (filled in by protocol module)
self.uplink = None
self.start_ts = int(time.time())
# UID generators, for servers that need it
self.uidgen = {}
def __init__(self, netname, proto):
# Initialize some variables
self.name = netname.lower()
self.conf = conf
self.serverdata = conf['servers'][netname]
self.sid = self.serverdata["sid"]
self.botdata = conf['bot']
self.proto = proto
self.pingfreq = self.serverdata.get('pingfreq') or 30
self.pingtimeout = self.pingfreq * 2
self.connected = threading.Event()
self.aborted = threading.Event()
self.initVars()
if world.testing:
# HACK: Don't thread if we're running tests.
self.connect()
else:
self.connection_thread = threading.Thread(target = self.connect)
self.connection_thread.start()
self.pingTimer = None
def connect(self):
ip = self.serverdata["ip"]
port = self.serverdata["port"]
while True:
self.initVars()
checks_ok = True
try:
self.socket = socket.socket()
self.socket.setblocking(0)
# Initial connection timeout is a lot smaller than the timeout after
# we've connected; this is intentional.
self.socket.settimeout(self.pingfreq)
self.ssl = self.serverdata.get('ssl')
if self.ssl:
log.info('(%s) Attempting SSL for this connection...', self.name)
certfile = self.serverdata.get('ssl_certfile')
keyfile = self.serverdata.get('ssl_keyfile')
if certfile and keyfile:
try:
self.socket = ssl.wrap_socket(self.socket,
certfile=certfile,
keyfile=keyfile)
except OSError:
log.exception('(%s) Caught OSError trying to '
'initialize the SSL connection; '
'are "ssl_certfile" and '
'"ssl_keyfile" set correctly?',
self.name)
checks_ok = False
else:
log.error('(%s) SSL certfile/keyfile was not set '
'correctly, aborting... ', self.name)
checks_ok = False
log.info("Connecting to network %r on %s:%s", self.name, ip, port)
self.socket.connect((ip, port))
self.socket.settimeout(self.pingtimeout)
if self.ssl and checks_ok:
peercert = self.socket.getpeercert(binary_form=True)
sha1fp = hashlib.sha1(peercert).hexdigest()
expected_fp = self.serverdata.get('ssl_fingerprint')
if expected_fp:
if sha1fp != expected_fp:
log.error('(%s) Uplink\'s SSL certificate '
'fingerprint (SHA1) does not match the '
'one configured: expected %r, got %r; '
'disconnecting...', self.name,
expected_fp, sha1fp)
checks_ok = False
else:
log.info('(%s) Uplink SSL certificate fingerprint '
'(SHA1) verified: %r', self.name, sha1fp)
else:
log.info('(%s) Uplink\'s SSL certificate fingerprint '
'is %r. You can enhance the security of your '
'link by specifying this in a "ssl_fingerprint"'
' option in your server block.', self.name,
sha1fp)
if checks_ok:
self.proto.connect(self)
self.spawnMain()
log.info('(%s) Starting ping schedulers....', self.name)
self.schedulePing()
log.info('(%s) Server ready; listening for data.', self.name)
self.run()
else:
log.error('(%s) A configuration error was encountered '
'trying to set up this connection. Please check'
' your configuration file and try again.',
self.name)
except (socket.error, ProtocolError, ConnectionError) as e:
log.warning('(%s) Disconnected from IRC: %s: %s',
self.name, type(e).__name__, str(e))
self._disconnect()
autoconnect = self.serverdata.get('autoconnect')
log.debug('(%s) Autoconnect delay set to %s seconds.', self.name, autoconnect)
if autoconnect is not None and autoconnect >= 0:
log.info('(%s) Going to auto-reconnect in %s seconds.', self.name, autoconnect)
time.sleep(autoconnect)
else:
return
def _disconnect(self):
log.debug('(%s) Canceling pingTimer at %s due to _disconnect() call', self.name, time.time())
self.connected.clear()
try:
self.socket.close()
self.pingTimer.cancel()
except: # Socket timed out during creation; ignore
pass
# Internal hook signifying that a network has disconnected.
self.callHooks([None, 'PYLINK_DISCONNECT', {}])
def disconnect(self):
"""Closes the IRC connection."""
self.aborted.set()
def run(self):
buf = b""
data = b""
while not self.aborted.is_set():
data = self.socket.recv(2048)
buf += data
if self.connected.is_set() and not data:
log.warning('(%s) No data received and self.connected is set; disconnecting!', self.name)
return
elif (time.time() - self.lastping) > self.pingtimeout:
log.warning('(%s) Connection timed out.', self.name)
return
while b'\n' in buf:
line, buf = buf.split(b'\n', 1)
line = line.strip(b'\r')
# FIXME: respect other encodings?
line = line.decode("utf-8", "replace")
log.debug("(%s) <- %s", self.name, line)
hook_args = None
try:
hook_args = self.proto.handle_events(self, line)
except Exception:
log.exception('(%s) Caught error in handle_events, disconnecting!', self.name)
return
# Only call our hooks if there's data to process. Handlers that support
# hooks will return a dict of parsed arguments, which can be passed on
# to plugins and the like. For example, the JOIN handler will return
# something like: {'channel': '#whatever', 'users': ['UID1', 'UID2',
# 'UID3']}, etc.
if hook_args is not None:
self.callHooks(hook_args)
def callHooks(self, hook_args):
numeric, command, parsed_args = hook_args
# Always make sure TS is sent.
if 'ts' not in parsed_args:
parsed_args['ts'] = int(time.time())
hook_cmd = command
hook_map = self.proto.hook_map
# Handlers can return a 'parse_as' key to send their payload to a
# different hook. An example of this is "/join 0" being interpreted
# as leaving all channels (PART).
if command in hook_map:
hook_cmd = hook_map[command]
hook_cmd = parsed_args.get('parse_as') or hook_cmd
log.debug('Parsed args %r received from %s handler (calling hook %s)', parsed_args, command, hook_cmd)
# Iterate over hooked functions, catching errors accordingly
for hook_func in world.command_hooks[hook_cmd]:
try:
log.debug('Calling function %s', hook_func)
hook_func(self, numeric, command, parsed_args)
except Exception:
# We don't want plugins to crash our servers...
log.exception('Unhandled exception caught in %r' % hook_func)
continue
def send(self, data):
# Safeguard against newlines in input!! Otherwise, each line gets
# treated as a separate command, which is particularly nasty.
data = data.replace('\n', ' ')
data = data.encode("utf-8") + b"\n"
stripped_data = data.decode("utf-8").strip("\n")
log.debug("(%s) -> %s", self.name, stripped_data)
try:
self.socket.send(data)
except (OSError, AttributeError):
log.debug("(%s) Dropping message %r; network isn't connected!", self.name, stripped_data)
def schedulePing(self):
self.proto.pingServer(self)
self.pingTimer = threading.Timer(self.pingfreq, self.schedulePing)
self.pingTimer.daemon = True
self.pingTimer.start()
log.debug('(%s) Ping scheduled at %s', self.name, time.time())
def spawnMain(self):
nick = self.botdata.get('nick') or 'PyLink'
ident = self.botdata.get('ident') or 'pylink'
host = self.serverdata["hostname"]
log.info('(%s) Connected! Spawning main client %s.', self.name, nick)
olduserobj = self.pseudoclient
self.pseudoclient = self.proto.spawnClient(self, nick, ident, host, modes={("+o", None)})
for chan in self.serverdata['channels']:
self.proto.joinClient(self, self.pseudoclient.uid, chan)
# PyLink internal hook called when spawnMain is called and the
# contents of Irc().pseudoclient change.
self.callHooks([self.sid, 'PYLINK_SPAWNMAIN', {'olduser': olduserobj}])
class IrcUser():
def __init__(self, nick, ts, uid, ident='null', host='null',
realname='PyLink dummy client', realhost='null',
ip='0.0.0.0'):
self.nick = nick
self.ts = ts
self.uid = uid
self.ident = ident
self.host = host
self.realhost = realhost
self.ip = ip
self.realname = realname
self.modes = set()
self.identified = False
self.channels = set()
self.away = ''
def __repr__(self):
return repr(self.__dict__)
class IrcServer():
"""PyLink IRC Server class.
uplink: The SID of this IrcServer instance's uplink. This is set to None
for the main PyLink PseudoServer!
name: The name of the server.
internal: Whether the server is an internal PyLink PseudoServer.
"""
def __init__(self, uplink, name, internal=False):
self.uplink = uplink
self.users = set()
self.internal = internal
self.name = name.lower()
def __repr__(self):
return repr(self.__dict__)
class IrcChannel():
def __init__(self):
self.users = set()
self.modes = {('n', None), ('t', None)}
self.topic = ''
self.ts = int(time.time())
self.topicset = False
self.prefixmodes = {'ops': set(), 'halfops': set(), 'voices': set(),
'owners': set(), 'admins': set()}
def __repr__(self):
return repr(self.__dict__)
def removeuser(self, target):
for s in self.prefixmodes.values():
s.discard(target)
self.users.discard(target)
### FakeIRC classes, used for test cases
class FakeIRC(Irc):
def connect(self):
self.messages = []
self.hookargs = []
self.hookmsgs = []
self.socket = None
self.initVars()
self.spawnMain()
self.connected = threading.Event()
self.connected.set()
def run(self, data):
"""Queues a message to the fake IRC server."""
log.debug('<- ' + data)
hook_args = self.proto.handle_events(self, data)
if hook_args is not None:
self.hookmsgs.append(hook_args)
self.callHooks(hook_args)
def send(self, data):
self.messages.append(data)
log.debug('-> ' + data)
def takeMsgs(self):
"""Returns a list of messages sent by the protocol module since
the last takeMsgs() call, so we can track what has been sent."""
msgs = self.messages
self.messages = []
return msgs
def takeCommands(self, msgs):
"""Returns a list of commands parsed from the output of takeMsgs()."""
sidprefix = ':' + self.sid
commands = []
for m in msgs:
args = m.split()
if m.startswith(sidprefix):
commands.append(args[1])
else:
commands.append(args[0])
return commands
def takeHooks(self):
"""Returns a list of hook arguments sent by the protocol module since
the last takeHooks() call."""
hookmsgs = self.hookmsgs
self.hookmsgs = []
return hookmsgs
class FakeProto():
"""Dummy protocol module for testing purposes."""
def __init__(self):
self.hook_map = {}
self.casemapping = 'rfc1459'
self.__name__ = 'FakeProto'
@staticmethod
def handle_events(irc, data):
pass
@staticmethod
def connect(irc):
pass
@staticmethod
def spawnClient(irc, nick, *args, **kwargs):
uid = randint(1, 10000000000)
ts = int(time.time())
irc.users[uid] = user = IrcUser(nick, ts, uid)
return user
@staticmethod
def joinClient(irc, client, channel):
irc.channels[channel].users.add(client)
irc.users[client].channels.add(channel)