modlib.py

   1 # Erebus IRC bot - Author: John Runyon
   2 # vim: fileencoding=utf-8
   3 # module helper functions, see modules/modtest.py for usage
   4 # This file is released into the public domain; see http://unlicense.org/
   5
   6 import sys
   7 import socket
   8 from functools import wraps
   9
  10 if sys.version_info.major < 3:
  11         stringbase = basestring
  12 else:
  13         stringbase = str
  14
  15 """Used to return an error to the bot core."""
  16 class error(object):
  17         def __init__(self, desc):
  18                 self.errormsg = desc
  19         def __nonzero__(self):
  20                 return False #object will test to False
  21         __bool__ = __nonzero__ #py3 compat
  22         def __repr__(self):
  23                 return '<modlib.error %r>' % self.errormsg
  24         def __str__(self):
  25                 return str(self.errormsg)
  26
  27 class modlib(object):
  28         # default (global) access levels
  29         OWNER   = 100
  30         MANAGER =  99
  31         ADMIN   =  75
  32         STAFF   =  50
  33         KNOWN   =   1
  34         AUTHED  =   0 # Users which have are known to be authed
  35         ANYONE  =  -1 # non-authed users have glevel set to -1
  36         IGNORED =  -2 # If the user is IGNORED, no hooks or chanhooks will fire for their messages. numhooks can still be fired.
  37         glevs   = {
  38                 'OWNER': OWNER,
  39                 'MANAGER': MANAGER,
  40                 'ADMIN': ADMIN,
  41                 'STAFF': STAFF,
  42                 'KNOWN': KNOWN,
  43                 'AUTHED': AUTHED,
  44                 'ANYONE': ANYONE,
  45                 'IGNORED': IGNORED,
  46         }
  47
  48         # (channel) access levels
  49         COWNER  =   5
  50         MASTER  =   4
  51         OP      =   3
  52         VOICE   =   2
  53         FRIEND  =   1
  54         PUBLIC  =   0 # Anyone (use glevel to control whether auth is needed)
  55         BANNED  =  -1 # The default reqclevel is PUBLIC, so any commands which needchan will be ignored from BANNED users unless the command reqclevel=-1
  56         #         [   0         1        2     3         4        5    -1]
  57         clevs   = [None, 'Friend', 'Voice', 'Op', 'Master', 'Owner', 'Banned']
  58
  59         # messages
  60         WRONGARGS = "Wrong number of arguments."
  61
  62         def __init__(self, name):
  63                 self.hooks = {} # {command:handler}
  64                 self.chanhooks = {} # {channel:handler}
  65                 self.exceptionhooks = [] # [(exception,handler)]
  66                 self.numhooks = {} # {numeric:handler}
  67                 self.sockhooks = [] # [(af,ty,address,handler_class)]
  68                 self.sockets = [] # [(sock,obj)]
  69                 self.helps = []
  70                 self.parent = None
  71
  72                 self.name = (name.split("."))[-1]
  73
  74         def modstart(self, parent):
  75                 #modstart can return a few things...
  76                 # None: unspecified success
  77                 # False: unspecified error
  78                 # modlib.error (or anything else False-y): specified error
  79                 # True: unspecified success
  80                 # non-empty string (or anything else True-y): specified success
  81                 #"specified" values will be printed. unspecified values will result in "OK" or "failed"
  82                 self.parent = parent
  83                 for cmd, func in self.hooks.items():
  84                         parent.hook(cmd, func)
  85                         parent.hook("%s.%s" % (self.name, cmd), func)
  86                 for chan, func in self.chanhooks.items():
  87                         parent.hookchan(chan, func)
  88                 for exc, func in self.exceptionhooks:
  89                         parent.hookexception(exc, func)
  90                 for num, func in self.numhooks.items():
  91                         parent.hooknum(num, func)
  92                 for hookdata in self.sockhooks:
  93                         self._create_socket(*hookdata)
  94
  95                 for func, args, kwargs in self.helps:
  96                         try:
  97                                 self.mod('help').reghelp(func, *args, **kwargs)
  98                         except:
  99                                 pass
 100                 return True
 101         def modstop(self, parent):
 102                 for cmd, func in self.hooks.items():
 103                         parent.unhook(cmd, func)
 104                         parent.unhook("%s.%s" % (self.name, cmd), func)
 105                 for chan, func in self.chanhooks.items():
 106                         parent.unhookchan(chan, func)
 107                 for exc, func in self.exceptionhooks:
 108                         parent.unhookexception(exc, func)
 109                 for num, func in self.numhooks.items():
 110                         parent.unhooknum(num, func)
 111                 for sock, obj in self.sockets:
 112                         self._destroy_socket(sock, obj)
 113
 114                 for func, args, kwargs in self.helps:
 115                         try:
 116                                 self.mod('help').dereghelp(func, *args, **kwargs)
 117                         except:
 118                                 pass
 119                 return True
 120
 121         def hook(self, cmd=None, needchan=True, glevel=ANYONE, clevel=PUBLIC, wantchan=None):
 122                 if wantchan is None: wantchan = needchan
 123                 _cmd = cmd #save this since it gets wiped out...
 124                 def realhook(func):
 125                         cmd = _cmd #...and restore it
 126                         if cmd is None:
 127                                 cmd = func.__name__ # default to function name
 128                         if isinstance(cmd, stringbase):
 129                                 cmd = (cmd,)
 130
 131                         if clevel > self.PUBLIC and not needchan:
 132                                 raise Exception('clevel must be left at default if needchan is False')
 133
 134                         func.needchan = needchan
 135                         func.wantchan = wantchan
 136                         func.reqglevel = glevel
 137                         func.reqclevel = clevel
 138                         func.cmd = cmd
 139                         func.module = func.__module__.split('.')[1]
 140
 141                         for c in cmd:
 142                                 self.hooks[c] = func
 143                                 if self.parent is not None:
 144                                         self.parent.hook(c, func)
 145                                         self.parent.hook("%s.%s" % (self.name, c), func)
 146                         return func
 147                 return realhook
 148
 149         def hookchan(self, chan, glevel=ANYONE, clevel=PUBLIC):
 150                 def realhook(func):
 151                         self.chanhooks[chan] = func
 152                         if self.parent is not None:
 153                                 self.parent.hookchan(chan, func)
 154                         return func
 155                 return realhook
 156
 157         def hookexception(self, exc):
 158                 def realhook(func):
 159                         self.exceptionhooks.append((exc, func))
 160                         if self.parent is not None:
 161                                 self.parent.hookexception(exc, func)
 162                         return func
 163                 return realhook
 164
 165         def hooknum(self, num):
 166                 def realhook(func):
 167                         self.numhooks[str(num)] = func
 168                         if self.parent is not None:
 169                                 self.parent.hooknum(str(num), func)
 170                         return func
 171                 return realhook
 172
 173         def bind(self, bindto, data=None):
 174                 """Used as a decorator on a class which implements getdata and parse methods.
 175                         See modules/sockets.py for an example.
 176                         Takes an arg like:
 177                         [unix:]/foo/bar
 178                         [udp|tcp:][ip:]port
 179                 """
 180                 if len(bindto) == 0:
 181                         raise Exception('bindto must have a value')
 182                 if bindto[0] == '/':
 183                         return self._hooksocket(socket.AF_UNIX, socket.SOCK_STREAM, bindto)
 184                 if len(bindto) > 5 and bindto[0:5] == 'unix:':
 185                         return self._hooksocket(socket.AF_UNIX, socket.SOCK_STREAM, bindto[5:])
 186                 af = socket.AF_INET
 187                 ty = socket.SOCK_STREAM
 188                 host = '0.0.0.0'
 189                 if len(bindto) > 4 and bindto[0:4] == 'udp:':
 190                         ty = socket.SOCK_DGRAM
 191                         bindto = bindto[4:]
 192                 if len(bindto) > 4 and bindto[0:4] == 'tcp:':
 193                         bindto = bindto[4:]
 194                 if ':' in bindto:
 195                         pieces = bindto.rsplit(':', 1)
 196                         host = pieces[0]
 197                         bindto = pieces[1]
 198                 port = int(bindto)
 199                 return self._hooksocket(af, ty, (host, port), data)
 200
 201         def bind_tcp(self, host, port, data=None):
 202                 return self._hooksocket(socket.AF_INET, socket.SOCK_STREAM, (host, port), data)
 203         def bind_udp(self, host, port, data=None):
 204                 return self._hooksocket(socket.AF_INET, socket.SOCK_DGRAM, (host, port), data)
 205         def bind_unix(self, path, data=None):
 206                 return self._hooksocket(socket.AF_UNIX, socket.SOCK_STREAM, path, data)
 207         def _hooksocket(self, af, ty, address, data):
 208                 def realhook(cls):
 209                         if not (hasattr(cls, 'getdata') and callable(cls.getdata)):
 210                                 # Check early that the object implements getdata.
 211                                 # If getdata ever returns a non-empty list, then a parse method must also exist, but we don't check that.
 212                                 raise Exception('Attempted to hook a socket without a class to process data')
 213                         self.sockhooks.append((af, ty, address, cls, data))
 214                         if self.parent is not None:
 215                                 self._create_socket(af, ty, address, cls, data)
 216                         return cls
 217                 return realhook
 218         def _create_socket(self, af, ty, address, cls, data):
 219                 ty = ty | socket.SOCK_NONBLOCK
 220                 sock = socket.socket(af, ty)
 221                 sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 222                 sock.bind(address)
 223                 obj = _ListenSocket(self, sock, cls, data)
 224                 self.sockets.append((sock,obj))
 225                 sock.listen(5)
 226                 self.parent.newfd(obj, sock.fileno())
 227                 self.parent.log(repr(obj), '?', 'Socket ready to accept new connections (%r, %r, %r, %r)' % (af, ty, address, cls))
 228         def _destroy_socket(self, sock, obj):
 229                 obj.close()
 230
 231         def mod(self, modname):
 232                 if self.parent is not None:
 233                         return self.parent.module(modname)
 234                 else:
 235                         return error('unknown parent')
 236
 237
 238         def flags(self, *flags):
 239                 """Parses out "flags" to a command, like `MODUNLOAD -AUTOLOAD somemodule`
 240                 @lib.hook()
 241                 @lib.flags('autoload', 'force')
 242                 def baz(bot, user, chan, realtarget, flags, *args)
 243
 244                 Note the added `flags` argument, which will be a dict - in this case `{'autounload':true,'force':false}`."""
 245                 def realhook(func):
 246                         func.flags = [f.lower() for f in flags]
 247
 248                         @wraps(func)
 249                         def parseargs(bot, user, chan, realtarget, *_args):
 250                                 args = list(_args) # we need a copy, also need a list, iterate over _args-tuple, mutate args-list
 251                                 found_flags = {f: False for f in flags}
 252                                 for arg in _args:
 253                                         if arg[0] == "-" and len(arg) > 1:
 254                                                 found_prefix = None
 255                                                 possible_flag = arg[1:].lower()
 256                                                 for f in flags:
 257                                                         if possible_flag == f: # Exact match?
 258                                                                 found_flags[possible_flag] = True
 259                                                                 args.remove(arg)
 260                                                                 found_prefix = None
 261                                                                 break
 262                                                         elif f.find(possible_flag) == 0: # Is the current word a prefix of a flag?
 263                                                                 if found_prefix is not None: # Is it also a prefix of another flag?
 264                                                                         return 'Error: %s is a prefix of multiple flags (%s, %s).' % (possible_flag, found_prefix[1], f)
 265                                                                 else:
 266                                                                         found_prefix = (arg,f)
 267                                                 if found_prefix is not None: # found (only one) prefix
 268                                                         found_flags[found_prefix[1]] = True
 269                                                         args.remove(found_prefix[0])
 270                                 return func(bot, user, chan, realtarget, found_flags, *args)
 271
 272                         return parseargs
 273                 return realhook
 274
 275
 276         def argsEQ(self, num):
 277                 def realhook(func):
 278                         @wraps(func)
 279                         def checkargs(bot, user, chan, realtarget, *args):
 280                                 adjuster = 0
 281                                 if hasattr(checkargs, 'flags'):
 282                                         adjuster = 1
 283                                 if len(args)-adjuster == num:
 284                                         return func(bot, user, chan, realtarget, *args)
 285                                 else:
 286                                         bot.msg(user, self.WRONGARGS)
 287                         return checkargs
 288                 return realhook
 289
 290         def argsGE(self, num):
 291                 def realhook(func):
 292                         @wraps(func)
 293                         def checkargs(bot, user, chan, realtarget, *args):
 294                                 adjuster = 0
 295                                 if hasattr(checkargs, 'flags'):
 296                                         adjuster = 1
 297                                 if len(args)-adjuster >= num:
 298                                         return func(bot, user, chan, realtarget, *args)
 299                                 else:
 300                                         bot.msg(user, self.WRONGARGS)
 301                         return checkargs
 302                 return realhook
 303
 304         def help(self, *args, **kwargs):
 305                 """help(syntax, shorthelp, longhelp?, more lines longhelp?, cmd=...?)
 306                 Example:
 307                 help("<user> <pass>", "login")
 308                         ^ Help will only be one line. Command name determined based on function name.
 309                 help("<user> <level>", "add a user", cmd=("adduser", "useradd"))
 310                         ^ Help will be listed under ADDUSER; USERADD will say "alias for adduser"
 311                 help(None, "do stuff", "This command is really complicated.")
 312                         ^ Command takes no args. Short description (in overall HELP listing) is "do stuff".
 313                         Long description (HELP <command>) will say "<command> - do stuff", newline, "This command is really complicated."
 314                 """
 315                 def realhook(func):
 316                         if self.parent is not None:
 317                                 try:
 318                                         self.mod('help').reghelp(func, *args, **kwargs)
 319                                 except:
 320                                         pass
 321                         self.helps.append((func, args, kwargs))
 322                         return func
 323                 return realhook
 324
 325 class _ListenSocket(object):
 326         def __init__(self, lib, sock, cls, data):
 327                 self.clients = []
 328                 self.lib = lib
 329                 self.sock = sock
 330                 self.cls = cls
 331                 self.data = data
 332
 333         def _make_closer(self, obj, client):
 334                 def close():
 335                         self.lib.parent.log(repr(self), '?', 'Closing child socket #%d' % (client.fileno()))
 336                         try:
 337                                 obj.closing()
 338                         except AttributeError:
 339                                 pass
 340                         self.lib.parent.delfd(client.fileno())
 341                         client.shutdown(socket.SHUT_RDWR)
 342                         client.close()
 343                         self.clients.remove((client,obj))
 344                 return close
 345
 346         def getdata(self):
 347                 client, addr = self.sock.accept()
 348                 obj = self.cls(client, self.data)
 349                 obj.close = self._make_closer(obj, client)
 350                 self.lib.parent.log(repr(self), '?', 'New connection #%d from %s' % (client.fileno(), addr))
 351                 self.clients.append((client,obj))
 352                 self.lib.parent.newfd(obj, client.fileno())
 353                 return []
 354
 355         def close(self):
 356                 self.lib.parent.log(repr(self), '?', 'Socket closing')
 357                 if self.sock.fileno() != -1:
 358                         self.lib.parent.delfd(self.sock.fileno())
 359                         self.sock.shutdown(socket.SHUT_RDWR)
 360                         self.sock.close()
 361                 for client, obj in self.clients:
 362                         obj.close()
 363
 364         def __repr__(self): return '<_ListenSocket #%d>' % (self.sock.fileno())