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