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