etc: Add environment option to run init.d/ashd silently.
[ashd.git] / python3 / ashd / util.py
1 """High-level utility module for ashd(7)
2
3 This module implements a rather convenient interface for writing ashd
4 handlers, wrapping the low-level ashd.proto module.
5 """
6
7 import os, socket, collections
8 from . import proto
9
10 __all__ = ["stdfork", "pchild", "respond", "serveloop"]
11
12 def stdfork(argv, chinit = None):
13     """Fork a persistent handler process using the `argv' argument
14     list, as per the standard ashd(7) calling convention. For an
15     easier-to-use interface, see the `pchild' class.
16
17     If a callable object of no arguments is provided in the `chinit'
18     argument, it will be called in the child process before exec()'ing
19     the handler program, and can be used to set parameters for the new
20     process, such as working directory, nice level or ulimits.
21
22     Returns the file descriptor of the socket for sending requests to
23     the new child.
24     """
25     csk, psk = socket.socketpair(socket.AF_UNIX, socket.SOCK_SEQPACKET)
26     pid = os.fork()
27     if pid == 0:
28         try:
29             os.dup2(csk.fileno(), 0)
30             for fd in range(3, 1024):
31                 try:
32                     os.close(fd)
33                 except:
34                     pass
35             if chinit is not None:
36                 chinit()
37             os.execvp(argv[0], argv)
38         finally:
39             os._exit(127)
40     csk.close()
41     fd = os.dup(psk.fileno())
42     psk.close()
43     return fd
44
45 class pchild(object):
46     """class pchild(argv, autorespawn=False, chinit=None)
47
48     Represents a persistent child handler process, started as per the
49     standard ashd(7) calling convention. It will be called with the
50     `argv' argument lest, which should be a list (or other iterable)
51     of strings.
52
53     If `autorespawn' is specified as True, the child process will be
54     automatically restarted if a request cannot be successfully sent
55     to it.
56
57     For a description of the `chinit' argument, see `stdfork'.
58
59     When this child handler should be disposed of, care should be
60     taken to call the close() method to release its socket and let it
61     exit. This class also implements the resource-manager interface,
62     so that it can be used in `with' statements.
63     """
64     
65     def __init__(self, argv, autorespawn = False, chinit = None):
66         self.argv = argv
67         self.chinit = chinit
68         self.fd = -1
69         self.respawn = autorespawn
70         self.spawn()
71
72     def spawn(self):
73         """Start the child handler, or restart it if it is already
74         running. You should not have to call this method manually
75         unless you explicitly want to manage the process' lifecycle.
76         """
77         self.close()
78         self.fd = stdfork(self.argv, self.chinit)
79
80     def close(self):
81         """Close this child handler's socket. For normal child
82         handlers, this will make the program terminate normally.
83         """
84         if self.fd >= 0:
85             os.close(self.fd)
86             self.fd = -1
87
88     def __del__(self):
89         self.close()
90
91     def passreq(self, req):
92         """Pass the specified request (which should be an instance of
93         the ashd.proto.req class) to this child handler. If the child
94         handler fails for some reason, and `autorespawn' was specified
95         as True when creating this handler, one attempt will be made
96         to restart it.
97
98         Note: You still need to close the request normally.
99
100         This method may raise an OSError if the request fails and
101         autorespawning was either not requested, or if the
102         autorespawning fails.
103         """
104         try:
105             proto.sendreq(self.fd, req)
106         except OSError:
107             if self.respawn:
108                 self.spawn()
109                 proto.sendreq(self.fd, req)
110
111     def __enter__(self):
112         return self
113
114     def __exit__(self, *excinfo):
115         self.close()
116         return False
117
118 def respond(req, body, status = ("200 OK"), ctype = "text/html"):
119     """Simple function for conveniently responding to a request.
120
121     Sends the specified body text to the request's response socket,
122     prepending an HTTP header with the appropriate Content-Type and
123     Content-Length headers, and then closes the response socket.
124
125     The `status' argument can be used to specify a non-200 response,
126     and the `ctype' argument can be used to specify a non-HTML
127     MIME-type.
128
129     If `body' is not a byte string, its string representation will be
130     encoded as UTF-8.
131
132     For example:
133         respond(req, "Not found", status = "404 Not Found", ctype = "text/plain")
134     """
135     if isinstance(body, collections.ByteString):
136         body = bytes(body)
137     else:
138         body = str(body)
139         body = body.encode("utf-8")
140         if ctype[:5] == "text/" and ctype.find(';') < 0:
141             ctype = ctype + "; charset=utf-8"
142     try:
143         head = ""
144         head += "HTTP/1.1 %s\n" % status
145         head += "Content-Type: %s\n" % ctype
146         head += "Content-Length: %i\n" % len(body)
147         head += "\n"
148         req.sk.write(head.encode("ascii"))
149         req.sk.write(body)
150     finally:
151         req.close()
152
153 def serveloop(handler, sock = 0):
154     """Implements a simple loop for serving requests sequentially, by
155     receiving requests from standard input (or the specified socket),
156     passing them to the specified handler function, and finally making
157     sure to close them. Returns when end-of-file is received on the
158     incoming socket.
159
160     The handler function should be a callable object of one argument,
161     and is called once for each received request.
162     """
163     while True:
164         try:
165             req = proto.recvreq(sock)
166         except InterruptedError:
167             continue
168         if req is None:
169             break
170         try:
171             handler(req)
172         finally:
173             req.close()