python: Added the beginning of a client library for ashd.pdm.
[ashd.git] / python / ashd / pdmc.py
1 """Management for daemon processes
2
3 This module provides some client support for the daemon management
4 provided in the ashd.pdm module.
5 """
6
7 import socket
8
9 class protoerr(Exception):
10     pass
11
12 def resolve(spec):
13     if isinstance(spec, socket.socket):
14         return spec
15     sk = None
16     try:
17         if "/" in spec:
18             sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
19             sk.connect(spec)
20         elif spec.isdigit():
21             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
22             sk.connect(("localhost", int(spec)))
23         elif ":" in spec:
24             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
25             p = spec.rindex(":")
26             sk.connect((spec[:p], int(spec[p + 1:])))
27         else:
28             raise Exception("Unknown target specification %r" % spec)
29         rv = sk
30         sk = None
31     finally:
32         if sk is not None: sk.close()
33     return rv
34
35 class client(object):
36     def __init__(self, sk, proto = None):
37         self.sk = resolve(sk)
38         self.buf = ""
39         line = self.readline()
40         if line != "+PDM1":
41             raise protoerr("Illegal protocol signature")
42         if proto is not None:
43             self.select(proto)
44
45     def close(self):
46         self.sk.close()
47
48     def readline(self):
49         while True:
50             p = self.buf.find("\n")
51             if p >= 0:
52                 ret = self.buf[:p]
53                 self.buf = self.buf[p + 1:]
54                 return ret
55             ret = self.sk.recv(1024)
56             if ret == "":
57                 return None
58             self.buf += ret
59
60     def select(self, proto):
61         if "\n" in proto:
62             raise Exception("Illegal protocol specified: %r" % proto)
63         self.sk.send(proto + "\n")
64         rep = self.readline()
65         if len(rep) < 1 or rep[0] != "+":
66             raise protoerr("Error reply when selecting protocol %s: %s" % (proto, rep[1:]))
67
68     def __enter__(self):
69         return self
70
71     def __exit__(self, *excinfo):
72         self.close()
73         return False
74
75 class replclient(client):
76     def __init__(self, sk):
77         super(replclient, self).__init__(sk, "repl")
78
79     def run(self, code):
80         while True:
81             ncode = code.replace("\n\n", "\n")
82             if ncode == code: break
83             code = ncode
84         while len(code) > 0 and code[-1] == "\n":
85             code = code[:-1]
86         self.sk.send(code + "\n\n")
87         buf = ""
88         while True:
89             ln = self.readline()
90             if ln[0] == " ":
91                 buf += ln[1:] + "\n"
92             elif ln[0] == "+":
93                 return buf
94             elif ln[0] == "-":
95                 raise protoerr("Error reply: %s" % ln[1:])
96             else:
97                 raise protoerr("Illegal reply: %s" % ln)