Parse the socket mode explicitly in octal.
[pdm.git] / pdm / srv.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 import pickle, struct
10
11 __all__ = ["listener", "unixlistener", "tcplistener", "listen"]
12
13 protocols = {}
14
15 class repl(object):
16     def __init__(self, cl):
17         self.cl = cl
18         self.mod = types.ModuleType("repl")
19         self.mod.echo = self.echo
20         self.printer = pprint.PrettyPrinter(indent = 4, depth = 6)
21         cl.send("+REPL\n")
22
23     def sendlines(self, text):
24         for line in text.split("\n"):
25             self.cl.send(" " + line + "\n")
26
27     def echo(self, ob):
28         self.sendlines(self.printer.pformat(ob))
29
30     def command(self, cmd):
31         try:
32             try:
33                 ccode = compile(cmd, "PDM Input", "eval")
34             except SyntaxError:
35                 ccode = compile(cmd, "PDM Input", "exec")
36                 exec(ccode, self.mod.__dict__)
37                 self.cl.send("+OK\n")
38             else:
39                 self.echo(eval(ccode, self.mod.__dict__))
40                 self.cl.send("+OK\n")
41         except:
42             for line in traceback.format_exception(*sys.exc_info()):
43                 self.cl.send(" " + line)
44             self.cl.send("+EXC\n")
45
46     def handle(self, buf):
47         p = buf.find("\n\n")
48         if p < 0:
49             return buf
50         cmd = buf[:p + 1]
51         self.command(cmd)
52         return buf[p + 2:]
53 protocols["repl"] = repl
54
55 class perf(object):
56     def __init__(self, cl):
57         self.cl = cl
58         self.odtab = {}
59         cl.send("+PERF1\n")
60         self.buf = ""
61         self.lock = threading.Lock()
62         self.subscribed = {}
63
64     def closed(self):
65         for id, recv in self.subscribed.items():
66             ob = self.odtab[id]
67             if ob is None: continue
68             ob, protos = ob
69             try:
70                 ob.unsubscribe(recv)
71             except: pass
72
73     def send(self, *args):
74         self.lock.acquire()
75         try:
76             buf = pickle.dumps(args)
77             buf = struct.pack(">l", len(buf)) + buf
78             self.cl.send(buf)
79         finally:
80             self.lock.release()
81
82     def bindob(self, id, ob):
83         if not hasattr(ob, "pdm_protocols"):
84             raise ValueError("Object does not support PDM introspection")
85         try:
86             proto = ob.pdm_protocols()
87         except Exception as exc:
88             raise ValueError("PDM introspection failed", exc)
89         self.odtab[id] = ob, proto
90         return proto
91
92     def bind(self, id, module, obnm):
93         resmod = sys.modules.get(module)
94         if resmod is None:
95             self.send("-", ImportError("No such module: %s" % module))
96             return
97         try:
98             ob = getattr(resmod, obnm)
99         except AttributeError:
100             self.send("-", AttributeError("No such object: %s" % obnm))
101             return
102         try:
103             proto = self.bindob(id, ob)
104         except Exception as exc:
105             self.send("-", exc)
106             return
107         self.send("+", proto)
108
109     def getob(self, id, proto):
110         ob = self.odtab.get(id)
111         if ob is None:
112             self.send("-", ValueError("No such bound ID: %r" % id))
113             return None
114         ob, protos = ob
115         if proto not in protos:
116             self.send("-", ValueError("Object does not support that protocol"))
117             return None
118         return ob
119
120     def lookup(self, tgtid, srcid, obnm):
121         src = self.getob(srcid, "dir")
122         if src is None:
123             return
124         try:
125             ob = src.lookup(obnm)
126         except KeyError as exc:
127             self.send("-", exc)
128             return
129         try:
130             proto = self.bindob(tgtid, ob)
131         except Exception as exc:
132             self.send("-", exc)
133             return
134         self.send("+", proto)
135
136     def unbind(self, id):
137         ob = self.odtab.get(id)
138         if ob is None:
139             self.send("-", KeyError("No such name bound: %r" % id))
140             return
141         ob, protos = ob
142         del self.odtab[id]
143         recv = self.subscribed.get(id)
144         if recv is not None:
145             ob.unsubscribe(recv)
146             del self.subscribed[id]
147         self.send("+")
148
149     def listdir(self, id):
150         ob = self.getob(id, "dir")
151         if ob is None:
152             return
153         self.send("+", ob.listdir())
154
155     def readattr(self, id):
156         ob = self.getob(id, "attr")
157         if ob is None:
158             return
159         try:
160             ret = ob.readattr()
161         except Exception as exc:
162             self.send("-", Exception("Could not read attribute"))
163             return
164         self.send("+", ret)
165
166     def attrinfo(self, id):
167         ob = self.getob(id, "attr")
168         if ob is None:
169             return
170         self.send("+", ob.attrinfo())
171
172     def invoke(self, id, method, args, kwargs):
173         ob = self.getob(id, "invoke")
174         if ob is None:
175             return
176         try:
177             self.send("+", ob.invoke(method, *args, **kwargs))
178         except Exception as exc:
179             self.send("-", exc)
180
181     def event(self, id, ob, ev):
182         self.send("*", id, ev)
183
184     def subscribe(self, id):
185         ob = self.getob(id, "event")
186         if ob is None:
187             return
188         if id in self.subscribed:
189             self.send("-", ValueError("Already subscribed"))
190         def recv(ev):
191             self.event(id, ob, ev)
192         ob.subscribe(recv)
193         self.subscribed[id] = recv
194         self.send("+")
195
196     def unsubscribe(self, id):
197         ob = self.getob(id, "event")
198         if ob is None:
199             return
200         recv = self.subscribed.get(id)
201         if recv is None:
202             self.send("-", ValueError("Not subscribed"))
203         ob.unsubscribe(recv)
204         del self.subscribed[id]
205         self.send("+")
206
207     def command(self, data):
208         cmd = data[0]
209         if cmd == "bind":
210             self.bind(*data[1:])
211         elif cmd == "unbind":
212             self.unbind(*data[1:])
213         elif cmd == "lookup":
214             self.lookup(*data[1:])
215         elif cmd == "ls":
216             self.listdir(*data[1:])
217         elif cmd == "readattr":
218             self.readattr(*data[1:])
219         elif cmd == "attrinfo":
220             self.attrinfo(*data[1:])
221         elif cmd == "invoke":
222             self.invoke(*data[1:])
223         elif cmd == "subs":
224             self.subscribe(*data[1:])
225         elif cmd == "unsubs":
226             self.unsubscribe(*data[1:])
227         else:
228             self.send("-", Exception("Unknown command: %r" % (cmd,)))
229
230     def handle(self, buf):
231         if len(buf) < 4:
232             return buf
233         dlen = struct.unpack(">l", buf[:4])[0]
234         if len(buf) < dlen + 4:
235             return buf
236         data = pickle.loads(buf[4:dlen + 4])
237         self.command(data)
238         return buf[dlen + 4:]
239         
240 protocols["perf"] = perf
241
242 class client(threading.Thread):
243     def __init__(self, sk):
244         super(client, self).__init__(name = "Management client")
245         self.setDaemon(True)
246         self.sk = sk
247         self.handler = self
248
249     def send(self, data):
250         return self.sk.send(data)
251
252     def choose(self, proto):
253         if proto in protocols:
254             self.handler = protocols[proto](self)
255         else:
256             self.send("-ERR Unknown protocol: %s\n" % proto)
257             raise Exception()
258
259     def handle(self, buf):
260         p = buf.find("\n")
261         if p >= 0:
262             proto = buf[:p]
263             buf = buf[p + 1:]
264             self.choose(proto)
265         return buf
266
267     def run(self):
268         try:
269             buf = ""
270             self.send("+PDM1\n")
271             while True:
272                 ret = self.sk.recv(1024)
273                 if ret == "":
274                     return
275                 buf += ret
276                 while True:
277                     try:
278                         nbuf = self.handler.handle(buf)
279                     except:
280                         return
281                     if nbuf == buf:
282                         break
283                     buf = nbuf
284         finally:
285             #for line in traceback.format_exception(*sys.exc_info()):
286             #    print line
287             try:
288                 self.sk.close()
289             finally:
290                 if hasattr(self.handler, "closed"):
291                     self.handler.closed()
292             
293
294 class listener(threading.Thread):
295     def __init__(self):
296         super(listener, self).__init__(name = "Management listener")
297         self.setDaemon(True)
298
299     def listen(self, sk):
300         self.running = True
301         while self.running:
302             rfd, wfd, efd = select.select([sk], [], [sk], 1)
303             for fd in rfd:
304                 if fd == sk:
305                     nsk, addr = sk.accept()
306                     self.accept(nsk, addr)
307
308     def stop(self):
309         self.running = False
310         self.join()
311
312     def accept(self, sk, addr):
313         cl = client(sk)
314         cl.start()
315
316 class unixlistener(listener):
317     def __init__(self, name, mode = 0o600, group = None):
318         super(unixlistener, self).__init__()
319         self.name = name
320         self.mode = mode
321         self.group = group
322
323     def run(self):
324         sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
325         ul = False
326         try:
327             if os.path.exists(self.name) and os.path.stat.S_ISSOCK(os.stat(self.name).st_mode):
328                 os.unlink(self.name)
329             sk.bind(self.name)
330             ul = True
331             os.chmod(self.name, self.mode)
332             if self.group is not None:
333                 os.chown(self.name, os.getuid(), grp.getgrnam(self.group).gr_gid)
334             sk.listen(16)
335             self.listen(sk)
336         finally:
337             sk.close()
338             if ul:
339                 os.unlink(self.name)
340
341 class tcplistener(listener):
342     def __init__(self, port, bindaddr = "127.0.0.1"):
343         super(tcplistener, self).__init__()
344         self.port = port
345         self.bindaddr = bindaddr
346
347     def run(self):
348         sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
349         try:
350             sk.bind((self.bindaddr, self.port))
351             sk.listen(16)
352             self.listen(sk)
353         finally:
354             sk.close()
355
356 def listen(spec):
357     if ":" in spec:
358         first = spec[:spec.index(":")]
359         last = spec[spec.rindex(":") + 1:]
360     else:
361         first = spec
362         last = spec
363     if "/" in first:
364         parts = spec.split(":")
365         mode = 0o600
366         group = None
367         if len(parts) > 1:
368             mode = int(parts[1], 8)
369         if len(parts) > 2:
370             group = parts[2]
371         ret = unixlistener(parts[0], mode = mode, group = group)
372         ret.start()
373         return ret
374     if last.isdigit():
375         p = spec.rindex(":")
376         host = spec[:p]
377         port = int(spec[p + 1:])
378         ret = tcplistener(port, bindaddr = host)
379         ret.start()
380         return ret
381     raise ValueError("Unparsable listener specification: %r" % spec)
382
383 import pdm.perf