Merge branch 'master' into python3
[pdm.git] / 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