python/ashd/pdmc.py

   1 """Management for daemon processes
   2
   3 This module provides some client support for the daemon management
   4 provided in the ashd.pdm module.
   5 """
   6
   7 import socket
   8
   9 class protoerr(Exception):
  10     pass
  11
  12 def resolve(spec):
  13     if isinstance(spec, socket.socket):
  14         return spec
  15     sk = None
  16     try:
  17         if "/" in spec:
  18             sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  19             sk.connect(spec)
  20         elif spec.isdigit():
  21             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  22             sk.connect(("localhost", int(spec)))
  23         elif ":" in spec:
  24             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  25             p = spec.rindex(":")
  26             sk.connect((spec[:p], int(spec[p + 1:])))
  27         else:
  28             raise Exception("Unknown target specification %r" % spec)
  29         rv = sk
  30         sk = None
  31     finally:
  32         if sk is not None: sk.close()
  33     return rv
  34
  35 class client(object):
  36     def __init__(self, sk, proto = None):
  37         self.sk = resolve(sk)
  38         self.buf = ""
  39         line = self.readline()
  40         if line != "+PDM1":
  41             raise protoerr("Illegal protocol signature")
  42         if proto is not None:
  43             self.select(proto)
  44
  45     def close(self):
  46         self.sk.close()
  47
  48     def readline(self):
  49         while True:
  50             p = self.buf.find("\n")
  51             if p >= 0:
  52                 ret = self.buf[:p]
  53                 self.buf = self.buf[p + 1:]
  54                 return ret
  55             ret = self.sk.recv(1024)
  56             if ret == "":
  57                 return None
  58             self.buf += ret
  59
  60     def select(self, proto):
  61         if "\n" in proto:
  62             raise Exception("Illegal protocol specified: %r" % proto)
  63         self.sk.send(proto + "\n")
  64         rep = self.readline()
  65         if len(rep) < 1 or rep[0] != "+":
  66             raise protoerr("Error reply when selecting protocol %s: %s" % (proto, rep[1:]))
  67
  68     def __enter__(self):
  69         return self
  70
  71     def __exit__(self, *excinfo):
  72         self.close()
  73         return False
  74
  75 class replclient(client):
  76     def __init__(self, sk):
  77         super(replclient, self).__init__(sk, "repl")
  78
  79     def run(self, code):
  80         while True:
  81             ncode = code.replace("\n\n", "\n")
  82             if ncode == code: break
  83             code = ncode
  84         while len(code) > 0 and code[-1] == "\n":
  85             code = code[:-1]
  86         self.sk.send(code + "\n\n")
  87         buf = ""
  88         while True:
  89             ln = self.readline()
  90             if ln[0] == " ":
  91                 buf += ln[1:] + "\n"
  92             elif ln[0] == "+":
  93                 return buf
  94             elif ln[0] == "-":
  95                 raise protoerr("Error reply: %s" % ln[1:])
  96             else:
  97                 raise protoerr("Illegal reply: %s" % ln)