1 | #!/usr/bin/python3 |
---|
2 | |
---|
3 | # Copyright 2019 Fiona Klute |
---|
4 | # |
---|
5 | # Licensed under the Apache License, Version 2.0 (the "License"); |
---|
6 | # you may not use this file except in compliance with the License. |
---|
7 | # You may obtain a copy of the License at |
---|
8 | # |
---|
9 | # http://www.apache.org/licenses/LICENSE-2.0 |
---|
10 | # |
---|
11 | # Unless required by applicable law or agreed to in writing, software |
---|
12 | # distributed under the License is distributed on an "AS IS" BASIS, |
---|
13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
---|
14 | # See the License for the specific language governing permissions and |
---|
15 | # limitations under the License. |
---|
16 | |
---|
17 | import socket |
---|
18 | import subprocess |
---|
19 | |
---|
20 | from http.client import HTTPConnection |
---|
21 | from multiprocessing import Process |
---|
22 | |
---|
23 | class HTTPSubprocessConnection(HTTPConnection): |
---|
24 | def __init__(self, command, host, port=None, |
---|
25 | output_filter=None, |
---|
26 | timeout=socket._GLOBAL_DEFAULT_TIMEOUT, |
---|
27 | blocksize=8192): |
---|
28 | super(HTTPSubprocessConnection, self).__init__(host, port, timeout, |
---|
29 | source_address=None, |
---|
30 | blocksize=blocksize) |
---|
31 | # "command" must be a list containing binary and command line |
---|
32 | # parameters |
---|
33 | self.command = command |
---|
34 | # This will be the subprocess reference when connected |
---|
35 | self._sproc = None |
---|
36 | # The subprocess return code is stored here on close() |
---|
37 | self.returncode = None |
---|
38 | # The set_tunnel method of the super class is not supported |
---|
39 | # (see exception doc) |
---|
40 | self.set_tunnel = None |
---|
41 | # This method will be run in a separate process and filter the |
---|
42 | # stdout of self._sproc. Its arguments are self._sproc.stdout |
---|
43 | # and the socket back to the HTTP connection (write-only). |
---|
44 | self._output_filter = output_filter |
---|
45 | # output filter process |
---|
46 | self._fproc = None |
---|
47 | |
---|
48 | def connect(self): |
---|
49 | s_local, s_remote = socket.socketpair(socket.AF_UNIX, |
---|
50 | socket.SOCK_STREAM) |
---|
51 | s_local.settimeout(self.timeout) |
---|
52 | |
---|
53 | # TODO: Maybe capture stderr? |
---|
54 | if self._output_filter: |
---|
55 | self._sproc = subprocess.Popen(self.command, stdout=subprocess.PIPE, |
---|
56 | stdin=s_remote, close_fds=True, |
---|
57 | bufsize=0) |
---|
58 | self._fproc = Process(target=self._output_filter, |
---|
59 | args=(self._sproc.stdout, s_remote)) |
---|
60 | self._fproc.start() |
---|
61 | else: |
---|
62 | self._sproc = subprocess.Popen(self.command, stdout=s_remote, |
---|
63 | stdin=s_remote, close_fds=True, |
---|
64 | bufsize=0) |
---|
65 | s_remote.close() |
---|
66 | self.sock = s_local |
---|
67 | |
---|
68 | def close(self): |
---|
69 | # close socket to subprocess for writing |
---|
70 | if self.sock: |
---|
71 | self.sock.shutdown(socket.SHUT_WR) |
---|
72 | |
---|
73 | # Wait for the process to stop, send SIGTERM/SIGKILL if |
---|
74 | # necessary |
---|
75 | if self._sproc: |
---|
76 | try: |
---|
77 | self.returncode = self._sproc.wait(self.timeout) |
---|
78 | except subprocess.TimeoutExpired: |
---|
79 | try: |
---|
80 | self._sproc.terminate() |
---|
81 | self.returncode = self._sproc.wait(self.timeout) |
---|
82 | except subprocess.TimeoutExpired: |
---|
83 | self._sproc.kill() |
---|
84 | self.returncode = self._sproc.wait(self.timeout) |
---|
85 | |
---|
86 | # filter process receives HUP on pipe when the subprocess |
---|
87 | # terminates |
---|
88 | if self._fproc: |
---|
89 | self._fproc.join() |
---|
90 | |
---|
91 | # close the connection in the super class, which also calls |
---|
92 | # self.sock.close() |
---|
93 | super().close() |
---|