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