1import re
2import os
3import subprocess
4from datetime import datetime
5from typing import Dict
6
7from urllib.parse import urlparse
8
9from .result import ExecResult
10
11
12def _get_path(x):
13    return x["path"]
14
15
16class Nghttp:
17
18    def __init__(self, path, connect_addr=None, tmp_dir="/tmp"):
19        self.NGHTTP = path
20        self.CONNECT_ADDR = connect_addr
21        self.TMP_DIR = tmp_dir
22
23    @staticmethod
24    def get_stream(streams, sid):
25        sid = int(sid)
26        if sid not in streams:
27            streams[sid] = {
28                    "id": sid,
29                    "header": {},
30                    "request": {
31                        "id": sid,
32                        "body": b''
33                    },
34                    "response": {
35                        "id": sid,
36                        "body": b''
37                    },
38                    "paddings": [],
39                    "promises": []
40            }
41        return streams[sid] if sid in streams else None
42
43    def run(self, urls, timeout, options):
44        return self._baserun(urls, timeout, options)
45
46    def complete_args(self, url, _timeout, options: [str]) -> [str]:
47        if not isinstance(url, list):
48            url = [url]
49        u = urlparse(url[0])
50        args = [self.NGHTTP]
51        if self.CONNECT_ADDR:
52            connect_host = self.CONNECT_ADDR
53            args.append("--header=host: %s:%s" % (u.hostname, u.port))
54        else:
55            connect_host = u.hostname
56        if options:
57            args.extend(options)
58        for xurl in url:
59            xu = urlparse(xurl)
60            nurl = "%s://%s:%s/%s" % (u.scheme, connect_host, xu.port, xu.path)
61            if xu.query:
62                nurl = "%s?%s" % (nurl, xu.query)
63            args.append(nurl)
64        return args
65
66    def _baserun(self, url, timeout, options):
67        return self._run(self.complete_args(url, timeout, options))
68
69    def parse_output(self, btext) -> Dict:
70        # getting meta data and response body out of nghttp's output
71        # is a bit tricky. Without '-v' we just get the body. With '-v' meta
72        # data and timings in both directions are listed.
73        # We rely on response :status: to be unique and on
74        # response body not starting with space.
75        # Something not good enough for general purpose, but for these tests.
76        output = {}
77        body = ''
78        streams = {}
79        skip_indents = True
80        # chunk output into lines. nghttp mixes text
81        # meta output with bytes from the response body.
82        lines = [l.decode() for l in btext.split(b'\n')]
83        for lidx, l in enumerate(lines):
84            if len(l) == 0:
85                body += '\n'
86                continue
87            m = re.match(r'(.*)\[.*] recv \(stream_id=(\d+)\) (\S+): (\S*)', l)
88            if m:
89                body += m.group(1)
90                s = self.get_stream(streams, m.group(2))
91                hname = m.group(3)
92                hval = m.group(4)
93                print("stream %d header %s: %s" % (s["id"], hname, hval))
94                header = s["header"]
95                if hname in header:
96                    header[hname] += ", %s" % hval
97                else:
98                    header[hname] = hval
99                body = ''
100                continue
101
102            m = re.match(r'(.*)\[.*] recv HEADERS frame <.* stream_id=(\d+)>', l)
103            if m:
104                body += m.group(1)
105                s = self.get_stream(streams, m.group(2))
106                if s:
107                    print("stream %d: recv %d header" % (s["id"], len(s["header"])))
108                    response = s["response"]
109                    hkey = "header"
110                    if "header" in response:
111                        h = response["header"]
112                        if ":status" in h and int(h[":status"]) >= 200:
113                            hkey = "trailer"
114                        else:
115                            prev = {
116                                "header": h
117                            }
118                            if "previous" in response:
119                                prev["previous"] = response["previous"]
120                            response["previous"] = prev
121                    response[hkey] = s["header"]
122                    s["header"] = {}
123                body = ''
124                continue
125
126            m = re.match(r'(.*)\[.*] recv DATA frame <length=(\d+), .*stream_id=(\d+)>', l)
127            if m:
128                body += m.group(1)
129                s = self.get_stream(streams, m.group(3))
130                blen = int(m.group(2))
131                if s:
132                    print("stream %d: %d DATA bytes added" % (s["id"], blen))
133                    padlen = 0
134                    if len(lines) > lidx + 2:
135                        mpad = re.match(r' +\(padlen=(\d+)\)', lines[lidx+2])
136                        if mpad:
137                            padlen = int(mpad.group(1))
138                    s["paddings"].append(padlen)
139                    blen -= padlen
140                    s["response"]["body"] += body[-blen:].encode()
141                body = ''
142                skip_indents = True
143                continue
144
145            m = re.match(r'(.*)\[.*] recv PUSH_PROMISE frame <.* stream_id=(\d+)>', l)
146            if m:
147                body += m.group(1)
148                s = self.get_stream(streams, m.group(2))
149                if s:
150                    # headers we have are request headers for the PUSHed stream
151                    # these have been received on the originating stream, the promised
152                    # stream id it mentioned in the following lines
153                    print("stream %d: %d PUSH_PROMISE header" % (s["id"], len(s["header"])))
154                    if len(lines) > lidx+2:
155                        m2 = re.match(r'\s+\(.*promised_stream_id=(\d+)\)', lines[lidx+2])
156                        if m2:
157                            s2 = self.get_stream(streams, m2.group(1))
158                            s2["request"]["header"] = s["header"]
159                            s["promises"].append(s2)
160                    s["header"] = {}
161                continue
162
163            m = re.match(r'(.*)\[.*] recv (\S+) frame <length=(\d+), .*stream_id=(\d+)>', l)
164            if m:
165                print("recv frame %s on stream %s" % (m.group(2), m.group(4)))
166                body += m.group(1)
167                skip_indents = True
168                continue
169
170            m = re.match(r'(.*)\[.*] send (\S+) frame <length=(\d+), .*stream_id=(\d+)>', l)
171            if m:
172                print("send frame %s on stream %s" % (m.group(2), m.group(4)))
173                body += m.group(1)
174                skip_indents = True
175                continue
176
177            if skip_indents and l.startswith('      '):
178                continue
179
180            if '[' != l[0]:
181                skip_indents = None
182                body += l + '\n'
183
184        # the main request is done on the lowest odd numbered id
185        main_stream = 99999999999
186        for sid in streams:
187            s = streams[sid]
188            if "header" in s["response"] and ":status" in s["response"]["header"]:
189                s["response"]["status"] = int(s["response"]["header"][":status"])
190            if (sid % 2) == 1 and sid < main_stream:
191                main_stream = sid
192
193        output["streams"] = streams
194        if main_stream in streams:
195            output["response"] = streams[main_stream]["response"]
196            output["paddings"] = streams[main_stream]["paddings"]
197        return output
198
199    def _raw(self, url, timeout, options):
200        args = ["-v"]
201        if options:
202            args.extend(options)
203        r = self._baserun(url, timeout, args)
204        if 0 == r.exit_code:
205            r.add_results(self.parse_output(r.outraw))
206        return r
207
208    def get(self, url, timeout=5, options=None):
209        return self._raw(url, timeout, options)
210
211    def assets(self, url, timeout=5, options=None):
212        if not options:
213            options = []
214        options.extend(["-ans"])
215        r = self._baserun(url, timeout, options)
216        assets = []
217        if 0 == r.exit_code:
218            lines = re.findall(r'[^\n]*\n', r.stdout, re.MULTILINE)
219            for lidx, l in enumerate(lines):
220                m = re.match(r'\s*(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+/(.*)', l)
221                if m:
222                    assets.append({
223                        "path": m.group(7),
224                        "status": int(m.group(5)),
225                        "size": m.group(6)
226                    })
227        assets.sort(key=_get_path)
228        r.add_assets(assets)
229        return r
230
231    def post_data(self, url, data, timeout=5, options=None):
232        reqbody = ("%s/nghttp.req.body" % self.TMP_DIR)
233        with open(reqbody, 'wb') as f:
234            f.write(data.encode('utf-8'))
235        if not options:
236            options = []
237        options.extend(["--data=%s" % reqbody])
238        return self._raw(url, timeout, options)
239
240    def post_name(self, url, name, timeout=5, options=None):
241        reqbody = ("%s/nghttp.req.body" % self.TMP_DIR)
242        with open(reqbody, 'w') as f:
243            f.write("--DSAJKcd9876\n")
244            f.write("Content-Disposition: form-data; name=\"value\"; filename=\"xxxxx\"\n")
245            f.write("Content-Type: text/plain\n")
246            f.write("\n%s\n" % name)
247            f.write("--DSAJKcd9876\n")
248        if not options:
249            options = []
250        options.extend(["--data=%s" % reqbody])
251        return self._raw(url, timeout, options)
252
253    def upload(self, url, fpath, timeout=5, options=None):
254        if not options:
255            options = []
256        options.extend(["--data=%s" % fpath])
257        return self._raw(url, timeout, options)
258
259    def upload_file(self, url, fpath, timeout=5, options=None):
260        fname = os.path.basename(fpath)
261        reqbody = ("%s/nghttp.req.body" % self.TMP_DIR)
262        with open(fpath, 'rb') as fin:
263            with open(reqbody, 'wb') as f:
264                f.write(("""--DSAJKcd9876
265Content-Disposition: form-data; name="xxx"; filename="xxxxx"
266Content-Type: text/plain
267
268testing mod_h2
269--DSAJKcd9876
270Content-Disposition: form-data; name="file"; filename="%s"
271Content-Type: application/octet-stream
272Content-Transfer-Encoding: binary
273
274""" % fname).encode('utf-8'))
275                f.write(fin.read())
276                f.write("""
277--DSAJKcd9876""".encode('utf-8'))
278        if not options:
279            options = []
280        options.extend([
281            "--data=%s" % reqbody,
282            "--expect-continue",
283            "-HContent-Type: multipart/form-data; boundary=DSAJKcd9876"])
284        return self._raw(url, timeout, options)
285
286    def _run(self, args) -> ExecResult:
287        print(("execute: %s" % " ".join(args)))
288        start = datetime.now()
289        p = subprocess.run(args, capture_output=True, text=False)
290        return ExecResult(args=args, exit_code=p.returncode,
291                          stdout=p.stdout, stderr=p.stderr,
292                          duration=datetime.now() - start)
293