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