get-with-headers.py
124 lines
| 3.2 KiB
| text/x-python
|
PythonLexer
/ tests / get-with-headers.py
|
r46434 | #!/usr/bin/env python3 | ||
|
r2532 | |||
|
r8447 | """This does HTTP GET requests given a host:port and path and returns | ||
|
r2532 | a subset of the headers plus the body of the result.""" | ||
|
r36594 | from __future__ import absolute_import | ||
|
r27296 | |||
|
r35798 | import argparse | ||
|
r27296 | import json | ||
import os | ||||
import sys | ||||
|
r7054 | |||
|
r29455 | from mercurial import ( | ||
|
r40190 | pycompat, | ||
|
r29455 | util, | ||
) | ||||
httplib = util.httplib | ||||
|
r7054 | try: | ||
|
r27296 | import msvcrt | ||
|
r43346 | |||
|
r7054 | msvcrt.setmode(sys.stdout.fileno(), os.O_BINARY) | ||
msvcrt.setmode(sys.stderr.fileno(), os.O_BINARY) | ||||
except ImportError: | ||||
pass | ||||
|
r36594 | stdout = getattr(sys.stdout, 'buffer', sys.stdout) | ||
|
r35798 | parser = argparse.ArgumentParser() | ||
parser.add_argument('--twice', action='store_true') | ||||
parser.add_argument('--headeronly', action='store_true') | ||||
parser.add_argument('--json', action='store_true') | ||||
parser.add_argument('--hgproto') | ||||
|
r43346 | parser.add_argument( | ||
'--requestheader', | ||||
nargs='*', | ||||
default=[], | ||||
help='Send an additional HTTP request header. Argument ' | ||||
'value is <header>=<value>', | ||||
) | ||||
parser.add_argument('--bodyfile', help='Write HTTP response body to a file') | ||||
|
r35798 | parser.add_argument('host') | ||
parser.add_argument('path') | ||||
parser.add_argument('show', nargs='*') | ||||
|
r12182 | |||
|
r35798 | args = parser.parse_args() | ||
twice = args.twice | ||||
headeronly = args.headeronly | ||||
formatjson = args.json | ||||
hgproto = args.hgproto | ||||
|
r35799 | requestheaders = args.requestheader | ||
|
r30764 | |||
|
r12182 | tag = None | ||
|
r43346 | |||
|
r12182 | def request(host, path, show): | ||
|
r17017 | assert not path.startswith('/'), path | ||
|
r12182 | global tag | ||
headers = {} | ||||
if tag: | ||||
headers['If-None-Match'] = tag | ||||
|
r30764 | if hgproto: | ||
headers['X-HgProto-1'] = hgproto | ||||
|
r5561 | |||
|
r35799 | for header in requestheaders: | ||
key, value = header.split('=', 1) | ||||
headers[key] = value | ||||
|
r12182 | conn = httplib.HTTPConnection(host) | ||
|
r17017 | conn.request("GET", '/' + path, None, headers) | ||
|
r12182 | response = conn.getresponse() | ||
|
r43346 | stdout.write( | ||
b'%d %s\n' % (response.status, response.reason.encode('ascii')) | ||||
) | ||||
|
r18380 | if show[:1] == ['-']: | ||
|
r43346 | show = sorted( | ||
h for h, v in response.getheaders() if h.lower() not in show | ||||
) | ||||
|
r12182 | for h in [h.lower() for h in show]: | ||
if response.getheader(h, None) is not None: | ||||
|
r43346 | stdout.write( | ||
b"%s: %s\n" | ||||
% (h.encode('ascii'), response.getheader(h).encode('ascii')) | ||||
) | ||||
|
r18400 | if not headeronly: | ||
|
r36594 | stdout.write(b'\n') | ||
|
r23409 | data = response.read() | ||
|
r24543 | |||
|
r35799 | if args.bodyfile: | ||
bodyfh = open(args.bodyfile, 'wb') | ||||
else: | ||||
|
r36594 | bodyfh = stdout | ||
|
r35799 | |||
|
r24543 | # Pretty print JSON. This also has the beneficial side-effect | ||
# of verifying emitted JSON is well-formed. | ||||
if formatjson: | ||||
# json.dumps() will print trailing newlines. Eliminate them | ||||
# to make tests easier to write. | ||||
|
r43697 | data = pycompat.json_loads(data) | ||
|
r24543 | lines = json.dumps(data, sort_keys=True, indent=2).splitlines() | ||
for line in lines: | ||||
|
r40190 | bodyfh.write(pycompat.sysbytes(line.rstrip())) | ||
|
r35799 | bodyfh.write(b'\n') | ||
|
r24543 | else: | ||
|
r35799 | bodyfh.write(data) | ||
if args.bodyfile: | ||||
bodyfh.close() | ||||
|
r12182 | |||
|
r31791 | if twice and response.getheader('ETag', None): | ||
tag = response.getheader('ETag') | ||||
|
r12182 | |||
return response.status | ||||
|
r43346 | |||
|
r35798 | status = request(args.host, args.path, args.show) | ||
|
r12182 | if twice: | ||
|
r35798 | status = request(args.host, args.path, args.show) | ||
|
r12182 | |||
if 200 <= status <= 305: | ||||
|
r5561 | sys.exit(0) | ||
sys.exit(1) | ||||