X-Git-Url: https://jfr.im/git/erebus.git/blobdiff_plain/867df3938aa231509618fafcbb7989329f352304..3cec5bdc665d310ac628e80b007dd2d71ae7d7bd:/modlib.py diff --git a/modlib.py b/modlib.py index 91e0584..11ca1fb 100644 --- a/modlib.py +++ b/modlib.py @@ -3,6 +3,7 @@ # module helper functions, see modules/modtest.py for usage # This file is released into the public domain; see http://unlicense.org/ +import abc import sys import socket from functools import wraps @@ -33,7 +34,7 @@ class modlib(object): KNOWN = 1 AUTHED = 0 # Users which have are known to be authed ANYONE = -1 # non-authed users have glevel set to -1 - IGNORED = -2 # The default reqglevel is ANYONE, so any commands will be ignored from IGNORED users unless the command reglevel=-2 + IGNORED = -2 # If the user is IGNORED, no hooks or chanhooks will fire for their messages. numhooks can still be fired. glevs = { 'OWNER': OWNER, 'MANAGER': MANAGER, @@ -60,6 +61,7 @@ class modlib(object): WRONGARGS = "Wrong number of arguments." def __init__(self, name): + self.Socketlike = Socketlike self.hooks = {} # {command:handler} self.chanhooks = {} # {channel:handler} self.exceptionhooks = [] # [(exception,handler)] @@ -170,33 +172,59 @@ class modlib(object): return func return realhook - def bind_tcp(self, host, port): - return self._hooksocket(socket.AF_INET, socket.SOCK_STREAM, (host, port)) - def bind_udp(self, host, port): - return self._hooksocket(socket.AF_INET, socket.SOCK_DGRAM, (host, port)) - def bind_unix(self, path): - return self._hooksocket(socket.AF_UNIX, socket.SOCK_STREAM, path) - def _hooksocket(self, af, ty, address): + def bind(self, bindto, data=None): + """Used as a decorator on a class which implements getdata and parse methods. + See modules/sockets.py for an example. + Takes an arg like: + [unix:]/foo/bar + [udp|tcp:][ip:]port + """ + if len(bindto) == 0: + raise Exception('bindto must have a value') + if bindto[0] == '/': + return self._hooksocket(socket.AF_UNIX, socket.SOCK_STREAM, bindto) + if len(bindto) > 5 and bindto[0:5] == 'unix:': + return self._hooksocket(socket.AF_UNIX, socket.SOCK_STREAM, bindto[5:]) + af = socket.AF_INET + ty = socket.SOCK_STREAM + host = '0.0.0.0' + if len(bindto) > 4 and bindto[0:4] == 'udp:': + ty = socket.SOCK_DGRAM + bindto = bindto[4:] + if len(bindto) > 4 and bindto[0:4] == 'tcp:': + bindto = bindto[4:] + if ':' in bindto: + pieces = bindto.rsplit(':', 1) + host = pieces[0] + bindto = pieces[1] + port = int(bindto) + return self._hooksocket(af, ty, (host, port), data) + + def bind_tcp(self, host, port, data=None): + return self._hooksocket(socket.AF_INET, socket.SOCK_STREAM, (host, port), data) + def bind_udp(self, host, port, data=None): + return self._hooksocket(socket.AF_INET, socket.SOCK_DGRAM, (host, port), data) + def bind_unix(self, path, data=None): + return self._hooksocket(socket.AF_UNIX, socket.SOCK_STREAM, path, data) + def _hooksocket(self, af, ty, address, data): def realhook(cls): - if not (hasattr(cls, 'getdata') and callable(cls.getdata)): - # Check early that the object implements getdata. - # If getdata ever returns a non-empty list, then a parse method must also exist, but we don't check that. + if not issubclass(cls, Socketlike): raise Exception('Attempted to hook a socket without a class to process data') - self.sockhooks.append((af, ty, address, cls)) + self.sockhooks.append((af, ty, address, cls, data)) if self.parent is not None: - self._create_socket(af, ty, address, cls) + self._create_socket(af, ty, address, cls, data) return cls return realhook - def _create_socket(self, af, ty, address, cls): + def _create_socket(self, af, ty, address, cls, data): ty = ty | socket.SOCK_NONBLOCK sock = socket.socket(af, ty) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind(address) - obj = _ListenSocket(self, sock, cls) + obj = _ListenSocket(self, sock, cls, data) self.sockets.append((sock,obj)) sock.listen(5) self.parent.newfd(obj, sock.fileno()) - self.parent.log(repr(obj), '?', 'Socket ready to accept new connections') + self.parent.log(repr(obj), '?', 'Socket ready to accept new connections (%r, %r, %r, %r)' % (af, ty, address, cls)) def _destroy_socket(self, sock, obj): obj.close() @@ -206,11 +234,53 @@ class modlib(object): else: return error('unknown parent') + + def flags(self, *flags): + """Parses out "flags" to a command, like `MODUNLOAD -AUTOLOAD somemodule` + @lib.hook() + @lib.flags('autoload', 'force') + def baz(bot, user, chan, realtarget, flags, *args) + + Note the added `flags` argument, which will be a dict - in this case `{'autounload':true,'force':false}`.""" + def realhook(func): + func.flags = [f.lower() for f in flags] + + @wraps(func) + def parseargs(bot, user, chan, realtarget, *_args): + args = list(_args) # we need a copy, also need a list, iterate over _args-tuple, mutate args-list + found_flags = {f: False for f in flags} + for arg in _args: + if arg[0] == "-" and len(arg) > 1: + found_prefix = None + possible_flag = arg[1:].lower() + for f in flags: + if possible_flag == f: # Exact match? + found_flags[possible_flag] = True + args.remove(arg) + found_prefix = None + break + elif f.find(possible_flag) == 0: # Is the current word a prefix of a flag? + if found_prefix is not None: # Is it also a prefix of another flag? + return 'Error: %s is a prefix of multiple flags (%s, %s).' % (possible_flag, found_prefix[1], f) + else: + found_prefix = (arg,f) + if found_prefix is not None: # found (only one) prefix + found_flags[found_prefix[1]] = True + args.remove(found_prefix[0]) + return func(bot, user, chan, realtarget, found_flags, *args) + + return parseargs + return realhook + + def argsEQ(self, num): def realhook(func): @wraps(func) def checkargs(bot, user, chan, realtarget, *args): - if len(args) == num: + adjuster = 0 + if hasattr(checkargs, 'flags'): + adjuster = 1 + if len(args)-adjuster == num: return func(bot, user, chan, realtarget, *args) else: bot.msg(user, self.WRONGARGS) @@ -221,7 +291,10 @@ class modlib(object): def realhook(func): @wraps(func) def checkargs(bot, user, chan, realtarget, *args): - if len(args) >= num: + adjuster = 0 + if hasattr(checkargs, 'flags'): + adjuster = 1 + if len(args)-adjuster >= num: return func(bot, user, chan, realtarget, *args) else: bot.msg(user, self.WRONGARGS) @@ -249,12 +322,58 @@ class modlib(object): return func return realhook -class _ListenSocket(object): - def __init__(self, lib, sock, cls): +class Socketlike(abc.ABC): + def __init__(self, sock, data): + """This default method saves the socket in self.sock and creates self.buffer for getdata(). The data is discarded.""" + self.sock = sock + self.buffer = b'' + + def getdata(self): + """This default method gets LF or CRLF separated lines from the socket and returns an array of completely-seen lines to the core. + This should work well for most line-based protocols (like IRC).""" + recvd = self.sock.recv(8192) + if recvd == b"": # EOF + if len(self.buffer) != 0: + # Process what's left in the buffer. We'll get called again after. + remaining_buf = self.buffer.decode('utf-8', 'backslashreplace') + self.buffer = b"" + return [remaining_buf] + else: + # Nothing left in the buffer. Return None to signal the core to close this socket. + return None + self.buffer += recvd + lines = [] + + while b"\n" in self.buffer: + pieces = self.buffer.split(b"\n", 1) + s = pieces[0].decode('utf-8', 'backslashreplace').rstrip("\r") + lines.append(pieces[0].decode('utf-8', 'backslashreplace')) + self.buffer = pieces[1] + + return lines + + def __str__(self): + return '%s#%d' % (self.__class__.__name__, self.sock.fileno()) + def __repr__(self): + return '<%s.%s #%d %s:%d>' % ((self.__class__.__module__, self.__class__.__name__, self.sock.fileno())+self.sock.getpeername()) + + @abc.abstractmethod + def parse(self, chunk): pass + + @classmethod + def __subclasshook__(cls, C): + if cls is Socketlike: + if any('parse' in B.__dict__ and 'getdata' in B.__dict__ for B in C.__mro__): + return True + return NotImplemented + +class _ListenSocket(Socketlike): + def __init__(self, lib, sock, cls, data): self.clients = [] self.lib = lib self.sock = sock self.cls = cls + self.data = data def _make_closer(self, obj, client): def close(): @@ -269,9 +388,13 @@ class _ListenSocket(object): self.clients.remove((client,obj)) return close + def parse(self): + # getdata will never return a non-empty array, so parse will never be called; but Socketlike requires this method + pass + def getdata(self): client, addr = self.sock.accept() - obj = self.cls(client) + obj = self.cls(client, self.data) obj.close = self._make_closer(obj, client) self.lib.parent.log(repr(self), '?', 'New connection #%d from %s' % (client.fileno(), addr)) self.clients.append((client,obj))