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