python/ashd/pdm.py

   1 """Management for daemon processes
   2
   3 This module contains a utility to listen for management commands on a
   4 socket, lending itself to managing daemon processes.
   5 """
   6
   7 import os, sys, socket, threading, grp, select
   8 import types, pprint, traceback
   9
  10 class repl(object):
  11     def __init__(self, cl):
  12         self.cl = cl
  13         self.mod = types.ModuleType("repl")
  14         self.mod.echo = self.echo
  15         self.printer = pprint.PrettyPrinter(indent = 4, depth = 6)
  16         cl.sk.send("+REPL\n")
  17
  18     def sendlines(self, text):
  19         for line in text.split("\n"):
  20             self.cl.sk.send(" " + line + "\n")
  21
  22     def echo(self, ob):
  23         self.sendlines(self.printer.pformat(ob))
  24
  25     def command(self, cmd):
  26         try:
  27             try:
  28                 ccode = compile(cmd, "PDM Input", "eval")
  29             except SyntaxError:
  30                 ccode = compile(cmd, "PDM Input", "exec")
  31                 exec ccode in self.mod.__dict__
  32                 self.cl.sk.send("+OK\n")
  33             else:
  34                 self.echo(eval(ccode, self.mod.__dict__))
  35                 self.cl.sk.send("+OK\n")
  36         except:
  37             for line in traceback.format_exception(*sys.exc_info()):
  38                 self.cl.sk.send(" " + line)
  39             self.cl.sk.send("+EXC\n")
  40
  41     def handle(self, buf):
  42         p = buf.find("\n\n")
  43         if p < 0:
  44             return buf
  45         cmd = buf[:p + 1]
  46         self.command(cmd)
  47         return buf[p + 2:]
  48
  49 class client(threading.Thread):
  50     def __init__(self, sk):
  51         super(client, self).__init__(name = "Management client")
  52         self.setDaemon(True)
  53         self.sk = sk
  54         self.handler = self
  55
  56     def choose(self, proto):
  57         if proto == "repl":
  58             self.handler = repl(self)
  59         else:
  60             self.sk.send("-ERR Unknown protocol: %s\n" % proto)
  61             raise Exception()
  62
  63     def handle(self, buf):
  64         p = buf.find("\n")
  65         if p >= 0:
  66             proto = buf[:p]
  67             buf = buf[p + 1:]
  68             self.choose(proto)
  69         return buf
  70
  71     def run(self):
  72         try:
  73             buf = ""
  74             self.sk.send("+PDM1\n")
  75             while True:
  76                 ret = self.sk.recv(1024)
  77                 if ret == "":
  78                     return
  79                 buf += ret
  80                 while True:
  81                     try:
  82                         nbuf = self.handler.handle(buf)
  83                     except:
  84                         return
  85                     if nbuf == buf:
  86                         break
  87                     buf = nbuf
  88         finally:
  89             self.sk.close()
  90
  91 class listener(threading.Thread):
  92     def __init__(self):
  93         super(listener, self).__init__(name = "Management listener")
  94         self.setDaemon(True)
  95
  96     def listen(self, sk):
  97         self.running = True
  98         while self.running:
  99             rfd, wfd, efd = select.select([sk], [], [sk], 1)
 100             for fd in rfd:
 101                 if fd == sk:
 102                     nsk, addr = sk.accept()
 103                     self.accept(nsk, addr)
 104
 105     def stop(self):
 106         self.running = False
 107         self.join()
 108
 109     def accept(self, sk, addr):
 110         cl = client(sk)
 111         cl.start()
 112
 113 class unixlistener(listener):
 114     def __init__(self, name, mode = 0600, group = None):
 115         super(unixlistener, self).__init__()
 116         self.name = name
 117         self.mode = mode
 118         self.group = group
 119
 120     def run(self):
 121         sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
 122         ul = False
 123         try:
 124             if os.path.exists(self.name) and os.path.stat.S_ISSOCK(os.stat(self.name).st_mode):
 125                 os.unlink(self.name)
 126             sk.bind(self.name)
 127             ul = True
 128             os.chmod(self.name, self.mode)
 129             if self.group is not None:
 130                 os.chown(self.name, os.getuid(), grp.getgrnam(self.group).gr_gid)
 131             sk.listen(16)
 132             self.listen(sk)
 133         finally:
 134             sk.close()
 135             if ul:
 136                 os.unlink(self.name)
 137
 138 class tcplistener(listener):
 139     def __init__(self, port, bindaddr = "127.0.0.1"):
 140         super(tcplistener, self).__init__()
 141         self.port = port
 142         self.bindaddr = bindaddr
 143
 144     def run(self):
 145         sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 146         try:
 147             sk.bind((self.bindaddr, self.port))
 148             sk.listen(16)
 149             self.listen(sk)
 150         finally:
 151             sk.close()