python: Added the beginning of a client library for ashd.pdm.
[ashd.git] / 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()