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                s = self.get_stream(streams, m.group(1))
90                hname = m.group(2)
91                hval = m.group(3)
92                print("stream %d header %s: %s" % (s["id"], hname, hval))
93                header = s["header"]
94                if hname in header:
95                    header[hname] += ", %s" % hval
96                else:
97                    header[hname] = hval
98                body = ''
99                continue
100
101            m = re.match(r'\[.*] recv HEADERS frame <.* stream_id=(\d+)>', l)
102            if m:
103                s = self.get_stream(streams, m.group(1))
104                if s:
105                    print("stream %d: recv %d header" % (s["id"], len(s["header"])))
106                    response = s["response"]
107                    hkey = "header"
108                    if "header" in response:
109                        h = response["header"]
110                        if ":status" in h and int(h[":status"]) >= 200:
111                            hkey = "trailer"
112                        else:
113                            prev = {
114                                "header": h
115                            }
116                            if "previous" in response:
117                                prev["previous"] = response["previous"]
118                            response["previous"] = prev
119                    response[hkey] = s["header"]
120                    s["header"] = {}
121                body = ''
122                continue
123
124            m = re.match(r'(.*)\[.*] recv DATA frame <length=(\d+), .*stream_id=(\d+)>', l)
125            if m:
126                s = self.get_stream(streams, m.group(3))
127                body += m.group(1)
128                blen = int(m.group(2))
129                if s:
130                    print("stream %d: %d DATA bytes added" % (s["id"], blen))
131                    padlen = 0
132                    if len(lines) > lidx + 2:
133                        mpad = re.match(r' +\(padlen=(\d+)\)', lines[lidx+2])
134                        if mpad:
135                            padlen = int(mpad.group(1))
136                    s["paddings"].append(padlen)
137                    blen -= padlen
138                    s["response"]["body"] += body[-blen:].encode()
139                body = ''
140                skip_indents = True
141                continue
142
143            m = re.match(r'\[.*] recv PUSH_PROMISE frame <.* stream_id=(\d+)>', l)
144            if m:
145                s = self.get_stream(streams, m.group(1))
146                if s:
147                    # headers we have are request headers for the PUSHed stream
148                    # these have been received on the originating stream, the promised
149                    # stream id it mentioned in the following lines
150                    print("stream %d: %d PUSH_PROMISE header" % (s["id"], len(s["header"])))
151                    if len(lines) > lidx+2:
152                        m2 = re.match(r'\s+\(.*promised_stream_id=(\d+)\)', lines[lidx+2])
153                        if m2:
154                            s2 = self.get_stream(streams, m2.group(1))
155                            s2["request"]["header"] = s["header"]
156                            s["promises"].append(s2)
157                    s["header"] = {}
158                continue
159
160            m = re.match(r'(.*)\[.*] recv (\S+) frame <length=(\d+), .*stream_id=(\d+)>', l)
161            if m:
162                print("recv frame %s on stream %s" % (m.group(2), m.group(4)))
163                body += m.group(1)
164                skip_indents = True
165                continue
166
167            m = re.match(r'(.*)\[.*] send (\S+) frame <length=(\d+), .*stream_id=(\d+)>', l)
168            if m:
169                print("send frame %s on stream %s" % (m.group(2), m.group(4)))
170                body += m.group(1)
171                skip_indents = True
172                continue
173
174            if skip_indents and l.startswith('      '):
175                continue
176
177            if '[' != l[0]:
178                skip_indents = None
179                body += l + '\n'
180
181        # the main request is done on the lowest odd numbered id
182        main_stream = 99999999999
183        for sid in streams:
184            s = streams[sid]
185            if "header" in s["response"] and ":status" in s["response"]["header"]:
186                s["response"]["status"] = int(s["response"]["header"][":status"])
187            if (sid % 2) == 1 and sid < main_stream:
188                main_stream = sid
189
190        output["streams"] = streams
191        if main_stream in streams:
192            output["response"] = streams[main_stream]["response"]
193            output["paddings"] = streams[main_stream]["paddings"]
194        return output
195
196    def _raw(self, url, timeout, options):
197        args = ["-v"]
198        if options:
199            args.extend(options)
200        r = self._baserun(url, timeout, args)
201        if 0 == r.exit_code:
202            r.add_results(self.parse_output(r.outraw))
203        return r
204
205    def get(self, url, timeout=5, options=None):
206        return self._raw(url, timeout, options)
207
208    def assets(self, url, timeout=5, options=None):
209        if not options:
210            options = []
211        options.extend(["-ans"])
212        r = self._baserun(url, timeout, options)
213        assets = []
214        if 0 == r.exit_code:
215            lines = re.findall(r'[^\n]*\n', r.stdout, re.MULTILINE)
216            for lidx, l in enumerate(lines):
217                m = re.match(r'\s*(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+/(.*)', l)
218                if m:
219                    assets.append({
220                        "path": m.group(7),
221                        "status": int(m.group(5)),
222                        "size": m.group(6)
223                    })
224        assets.sort(key=_get_path)
225        r.add_assets(assets)
226        return r
227
228    def post_data(self, url, data, timeout=5, options=None):
229        reqbody = ("%s/nghttp.req.body" % self.TMP_DIR)
230        with open(reqbody, 'wb') as f:
231            f.write(data.encode('utf-8'))
232        if not options:
233            options = []
234        options.extend(["--data=%s" % reqbody])
235        return self._raw(url, timeout, options)
236
237    def post_name(self, url, name, timeout=5, options=None):
238        reqbody = ("%s/nghttp.req.body" % self.TMP_DIR)
239        with open(reqbody, 'w') as f:
240            f.write("--DSAJKcd9876\n")
241            f.write("Content-Disposition: form-data; name=\"value\"; filename=\"xxxxx\"\n")
242            f.write("Content-Type: text/plain\n")
243            f.write("\n%s\n" % name)
244            f.write("--DSAJKcd9876\n")
245        if not options:
246            options = []
247        options.extend(["--data=%s" % reqbody])
248        return self._raw(url, timeout, options)
249
250    def upload(self, url, fpath, timeout=5, options=None):
251        if not options:
252            options = []
253        options.extend(["--data=%s" % fpath])
254        return self._raw(url, timeout, options)
255
256    def upload_file(self, url, fpath, timeout=5, options=None):
257        fname = os.path.basename(fpath)
258        reqbody = ("%s/nghttp.req.body" % self.TMP_DIR)
259        with open(fpath, 'rb') as fin:
260            with open(reqbody, 'wb') as f:
261                f.write(("""--DSAJKcd9876
262Content-Disposition: form-data; name="xxx"; filename="xxxxx"
263Content-Type: text/plain
264
265testing mod_h2
266--DSAJKcd9876
267Content-Disposition: form-data; name="file"; filename="%s"
268Content-Type: application/octet-stream
269Content-Transfer-Encoding: binary
270
271""" % fname).encode('utf-8'))
272                f.write(fin.read())
273                f.write("""
274--DSAJKcd9876""".encode('utf-8'))
275        if not options:
276            options = []
277        options.extend([
278            "--data=%s" % reqbody,
279            "--expect-continue",
280            "-HContent-Type: multipart/form-data; boundary=DSAJKcd9876"])
281        return self._raw(url, timeout, options)
282
283    def _run(self, args) -> ExecResult:
284        print(("execute: %s" % " ".join(args)))
285        start = datetime.now()
286        p = subprocess.run(args, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
287        return ExecResult(args=args, exit_code=p.returncode,
288                          stdout=p.stdout, stderr=p.stderr,
289                          duration=datetime.now() - start)
290