add ability to hook extra sockets

[erebus.git] / modlib.py

diff --git a/modlib.py b/modlib.py
index ddb201a252b9cc70e02ccc0d335f6eb105b2fc83..a04dd7878e22898f222cb06318ddcf6569b7e502 100644 (file)
--- a/modlib.py
+++ b/modlib.py
@@ -4,6 +4,7 @@
 # This file is released into the public domain; see http://unlicense.org/
 
 import sys
+import socket
 from functools import wraps
 
 if sys.version_info.major < 3:
@@ -59,10 +60,12 @@ class modlib(object):
        WRONGARGS = "Wrong number of arguments."
 
        def __init__(self, name):
-               self.hooks = {}
-               self.chanhooks = {}
-               self.exceptionhooks = []
-               self.numhooks = {}
+               self.hooks = {} # {command:handler}
+               self.chanhooks = {} # {channel:handler}
+               self.exceptionhooks = [] # [(exception,handler)]
+               self.numhooks = {} # {numeric:handler}
+               self.sockhooks = [] # [(af,ty,address,handler_class)]
+               self.sockets = [] # [(sock,obj)]
                self.helps = []
                self.parent = None
 
@@ -86,6 +89,8 @@ class modlib(object):
                        parent.hookexception(exc, func)
                for num, func in self.numhooks.items():
                        parent.hooknum(num, func)
+               for hookdata in self.sockhooks:
+                       self._create_socket(*hookdata)
 
                for func, args, kwargs in self.helps:
                        try:
@@ -103,6 +108,8 @@ class modlib(object):
                        parent.unhookexception(exc, func)
                for num, func in self.numhooks.items():
                        parent.unhooknum(num, func)
+               for sock, obj in self.sockets:
+                       self._destroy_socket(sock, obj)
 
                for func, args, kwargs in self.helps:
                        try:
@@ -163,6 +170,36 @@ 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 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.
+                               raise Exception('Attempted to hook a socket without a class to process data')
+                       self.sockhooks.append((af, ty, address, cls))
+                       if self.parent is not None:
+                               self._create_socket(af, ty, address, cls)
+                       return cls
+               return realhook
+       def _create_socket(self, af, ty, address, cls):
+               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)
+               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')
+       def _destroy_socket(self, sock, obj):
+               obj.close()
+
        def mod(self, modname):
                if self.parent is not None:
                        return self.parent.module(modname)
@@ -211,3 +248,44 @@ class modlib(object):
                        self.helps.append((func, args, kwargs))
                        return func
                return realhook
+
+class _ListenSocket(object):
+       def __init__(self, lib, sock, cls):
+               self.clients = []
+               self.lib = lib
+               self.sock = sock
+               self.cls = cls
+
+       def _make_closer(self, obj, client):
+               def close():
+                       print(repr(self), repr(obj))
+                       self.lib.parent.log(repr(self), '?', 'Closing child socket %d' % (client.fileno()))
+                       try:
+                               obj.closing()
+                       except AttributeError:
+                               pass
+                       self.lib.parent.delfd(client.fileno())
+                       client.shutdown(socket.SHUT_RDWR)
+                       client.close()
+                       self.clients.remove((client,obj))
+               return close
+
+       def getdata(self):
+               client, addr = self.sock.accept()
+               obj = self.cls(client)
+               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))
+               self.lib.parent.newfd(obj, client.fileno())
+               return []
+
+       def close(self):
+               self.lib.parent.log(repr(self), '?', 'Socket closing')
+               if self.sock.fileno() != -1:
+                       self.lib.parent.delfd(self.sock.fileno())
+                       self.sock.shutdown(socket.SHUT_RDWR)
+                       self.sock.close()
+               for client, obj in self.clients:
+                       obj.close()
+
+       def __repr__(self): return '<_ListenSocket #%d>' % (self.sock.fileno())