pdm/cli.py

   1 """Python Daemon Management -- Client functions
   2
   3 This module implements the client part of the PDM protocols. The
   4 primary objects of interest are the replclient and perfclient classes,
   5 which implement support for their respective protocols. See their
   6 documentation for details.
   7 """
   8
   9 import socket, pickle, struct, select, threading
  10
  11 __all__ = ["client", "replclient", "perfclient"]
  12
  13 class protoerr(Exception):
  14     """Raised on protocol errors"""
  15     pass
  16
  17 def resolve(spec):
  18     if isinstance(spec, socket.socket):
  19         return spec
  20     sk = None
  21     try:
  22         if "/" in spec:
  23             sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  24             sk.connect(spec)
  25         elif spec.isdigit():
  26             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  27             sk.connect(("localhost", int(spec)))
  28         elif ":" in spec:
  29             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  30             p = spec.rindex(":")
  31             sk.connect((spec[:p], int(spec[p + 1:])))
  32         else:
  33             raise Exception("Unknown target specification %r" % spec)
  34         rv = sk
  35         sk = None
  36     finally:
  37         if sk is not None: sk.close()
  38     return rv
  39
  40 class client(object):
  41     """PDM client
  42
  43     This class provides general facilities to speak to PDM servers,
  44     and is mainly intended to be subclassed to provide for the
  45     specific protocols, such as replclient and perfclient do.
  46
  47     `client' instances can be passed as arguments to select.select(),
  48     and can be used in `with' statements.
  49     """
  50     def __init__(self, sk, proto = None):
  51         """Create a client object connected to the specified
  52         server. `sk' can either be a socket object, which is used as
  53         it is, or a string specification very similar to the
  54         specification for L{pdm.srv.listen}, so see its documentation
  55         for details. The differences are only that this function does
  56         not take arguments specific to socket creation, like the mode
  57         and group arguments for Unix sockets. If `proto' is given,
  58         that subprotocol will negotiated with the server (by calling
  59         the select() method).
  60         """
  61         self.sk = resolve(sk)
  62         self.buf = ""
  63         line = self.readline()
  64         if line != "+PDM1":
  65             raise protoerr("Illegal protocol signature")
  66         if proto is not None:
  67             self.select(proto)
  68
  69     def close(self):
  70         """Close this connection"""
  71         self.sk.close()
  72
  73     def fileno(self):
  74         """Return the file descriptor of the underlying socket."""
  75         return self.sk.fileno()
  76
  77     def readline(self):
  78         """Read a single NL-terminated line and return it."""
  79         while True:
  80             p = self.buf.find("\n")
  81             if p >= 0:
  82                 ret = self.buf[:p]
  83                 self.buf = self.buf[p + 1:]
  84                 return ret
  85             ret = self.sk.recv(1024)
  86             if ret == "":
  87                 return None
  88             self.buf += ret
  89
  90     def select(self, proto):
  91         """Negotiate the given subprotocol with the server"""
  92         if "\n" in proto:
  93             raise Exception("Illegal protocol specified: %r" % proto)
  94         self.sk.send(proto + "\n")
  95         rep = self.readline()
  96         if len(rep) < 1 or rep[0] != "+":
  97             raise protoerr("Error reply when selecting protocol %s: %s" % (proto, rep[1:]))
  98
  99     def __enter__(self):
 100         return self
 101
 102     def __exit__(self, *excinfo):
 103         self.close()
 104         return False
 105
 106 class replclient(client):
 107     """REPL protocol client
 108
 109     Implements the client side of the REPL protocol; see
 110     L{pdm.srv.repl} for details on the protocol and its functionality.
 111     """
 112     def __init__(self, sk):
 113         """Create a connected client as documented in the `client' class."""
 114         super(replclient, self).__init__(sk, "repl")
 115
 116     def run(self, code):
 117         """Run a single block of Python code on the server. Returns
 118         the output of the command (as documented in L{pdm.srv.repl})
 119         as a string.
 120         """
 121         while True:
 122             ncode = code.replace("\n\n", "\n")
 123             if ncode == code: break
 124             code = ncode
 125         while len(code) > 0 and code[-1] == "\n":
 126             code = code[:-1]
 127         self.sk.send(code + "\n\n")
 128         buf = ""
 129         while True:
 130             ln = self.readline()
 131             if ln[0] == " ":
 132                 buf += ln[1:] + "\n"
 133             elif ln[0] == "+":
 134                 return buf
 135             elif ln[0] == "-":
 136                 raise protoerr("Error reply: %s" % ln[1:])
 137             else:
 138                 raise protoerr("Illegal reply: %s" % ln)
 139
 140 class perfproxy(object):
 141     def __init__(self, cl, id, proto):
 142         self.cl = cl
 143         self.id = id
 144         self.proto = proto
 145         self.subscribers = set()
 146
 147     def lookup(self, name):
 148         self.cl.lock.acquire()
 149         try:
 150             id = self.cl.nextid
 151             self.cl.nextid += 1
 152         finally:
 153             self.cl.lock.release()
 154         (proto,) = self.cl.run("lookup", id, self.id, name)
 155         proxy = perfproxy(self.cl, id, proto)
 156         self.cl.proxies[id] = proxy
 157         return proxy
 158
 159     def listdir(self):
 160         return self.cl.run("ls", self.id)[0]
 161
 162     def readattr(self):
 163         return self.cl.run("readattr", self.id)[0]
 164
 165     def attrinfo(self):
 166         return self.cl.run("attrinfo", self.id)[0]
 167
 168     def invoke(self, method, *args, **kwargs):
 169         return self.cl.run("invoke", self.id, method, args, kwargs)[0]
 170
 171     def subscribe(self, cb):
 172         if cb in self.subscribers:
 173             raise ValueError("Already subscribed")
 174         if len(self.subscribers) == 0:
 175             self.cl.run("subs", self.id)
 176         self.subscribers.add(cb)
 177
 178     def unsubscribe(self, cb):
 179         if cb not in self.subscribers:
 180             raise ValueError("Not subscribed")
 181         self.subscribers.remove(cb)
 182         if len(self.subscribers) == 0:
 183             self.cl.run("unsubs", self.id)
 184
 185     def notify(self, ev):
 186         for cb in self.subscribers:
 187             try:
 188                 cb(ev)
 189             except: pass
 190
 191     def close(self):
 192         if self.id is not None:
 193             self.cl.run("unbind", self.id)
 194             del self.cl.proxies[self.id]
 195             self.id = None
 196
 197     def __del__(self):
 198         self.close()
 199
 200     def __enter__(self):
 201         return self
 202
 203     def __exit__(self, *excinfo):
 204         self.close()
 205         return False
 206
 207 class perfclient(client):
 208     """PERF protocol client
 209
 210     Implements the client side of the PERF protocol; see
 211     L{pdm.srv.perf} for details on the protocol and its functionality.
 212
 213     This client class implements functions for finding PERF objects on
 214     the server, and returns, for each server-side object looked up, a
 215     proxy object that mimics exactly the PERF interfaces that the
 216     object implements. As the proxy objects reference live objects on
 217     the server, they should be released when they are no longer used;
 218     they implement a close() method for that purpose, and can also be
 219     used in `with' statements.
 220
 221     See L{pdm.srv.perf} for details on the various PERF interfaces
 222     that the proxy objects might implement.
 223     """
 224     def __init__(self, sk):
 225         """Create a connected client as documented in the `client' class."""
 226         super(perfclient, self).__init__(sk, "perf")
 227         self.nextid = 0
 228         self.lock = threading.Lock()
 229         self.proxies = {}
 230         self.names = {}
 231
 232     def send(self, ob):
 233         buf = pickle.dumps(ob)
 234         buf = struct.pack(">l", len(buf)) + buf
 235         self.sk.send(buf)
 236
 237     def recvb(self, num):
 238         buf = ""
 239         while len(buf) < num:
 240             data = self.sk.recv(num - len(buf))
 241             if data == "":
 242                 raise EOFError()
 243             buf += data
 244         return buf
 245
 246     def recv(self):
 247         return pickle.loads(self.recvb(struct.unpack(">l", self.recvb(4))[0]))
 248
 249     def event(self, id, ev):
 250         proxy = self.proxies.get(id)
 251         if proxy is None: return
 252         proxy.notify(ev)
 253
 254     def dispatch(self, timeout = None):
 255         """Wait for an incoming notification from the server, and
 256         dispatch it to the callback functions that have been
 257         registered for it. If `timeout' is specified, wait no longer
 258         than so many seconds; otherwise, wait forever. This client
 259         object may also be used as argument to select.select().
 260         """
 261         rfd, wfd, efd = select.select([self.sk], [], [], timeout)
 262         if self.sk in rfd:
 263             msg = self.recv()
 264             if msg[0] == "*":
 265                 self.event(msg[1], msg[2])
 266             else:
 267                 raise ValueError("Unexpected non-event message: %r" % msg[0])
 268
 269     def recvreply(self):
 270         while True:
 271             reply = self.recv()
 272             if reply[0] in ("+", "-"):
 273                 return reply
 274             elif reply[0] == "*":
 275                 self.event(reply[1], reply[2])
 276             else:
 277                 raise ValueError("Illegal reply header: %r" % reply[0])
 278
 279     def run(self, cmd, *args):
 280         self.lock.acquire()
 281         try:
 282             self.send((cmd,) + args)
 283             reply = self.recvreply()
 284             if reply[0] == "+":
 285                 return reply[1:]
 286             else:
 287                 raise reply[1]
 288         finally:
 289             self.lock.release()
 290
 291     def lookup(self, module, obnm):
 292         """Look up a single server-side object by the given name in
 293         the given module. Will return a new proxy object for each
 294         call when called multiple times for the same name.
 295         """
 296         self.lock.acquire()
 297         try:
 298             id = self.nextid
 299             self.nextid += 1
 300         finally:
 301             self.lock.release()
 302         (proto,) = self.run("bind", id, module, obnm)
 303         proxy = perfproxy(self, id, proto)
 304         self.proxies[id] = proxy
 305         return proxy
 306
 307     def find(self, name):
 308         """Convenience function for looking up server-side objects
 309         through PERF directories and for multiple uses. The object
 310         name can be given as "MODULE.OBJECT", which will look up the
 311         named OBJECT in the named MODULE, and can be followed by any
 312         number of slash-separated names, which will assume that the
 313         object to the left of the slash is a PERF directory, and will
 314         return the object in that directory by the name to the right
 315         of the slash. For instance, find("pdm.perf.sysres/cputime")
 316         will return the built-in attribute for reading the CPU time
 317         used by the server process.
 318
 319         The proxy objects returned by this function are cached and the
 320         same object are returned the next time the same name is
 321         requested, which means that they are kept live until the
 322         client connection is closed.
 323         """
 324         ret = self.names.get(name)
 325         if ret is None:
 326             if "/" in name:
 327                 p = name.rindex("/")
 328                 ret = self.find(name[:p]).lookup(name[p + 1:])
 329             else:
 330                 p = name.rindex(".")
 331                 ret = self.lookup(name[:p], name[p + 1:])
 332             self.names[name] = ret
 333         return ret