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