Easily communicate with JSON APIs

restest.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# restest.py - Easily communicate with JSON APIs.
#

import sys
from argparse import ArgumentParser
from urllib import urlencode
from urllib2 import build_opener, HTTPError

from pygments import highlight
from pygments.lexers import JsonLexer
from pygments.formatters import TerminalFormatter


COLORS = {
    'green': '\033[32m',
    'bgreen': '\033[1;32m',
    'bgrey': '\033[1;30m',
    'reset': '\033[0m'
}


def extract_values(values):
    exploded = map(lambda s: s.split('='), values)
    return dict(exploded)


ap = ArgumentParser(description="Easily communicate with JSON APIs")
ap.add_argument(
    "url",
    help="API endpoint to hit",
    metavar="URL"
)
ap.add_argument(
    "--post",
    help="POST data to send with request",
    nargs="*",
    metavar="DATA"
)
ap.add_argument(
    "--headers",
    help="Headers to add to the request",
    nargs="*",
    metavar="HEADER"
)
ap.add_argument(
    "--http-user",
    help="HTTP basic auth username",
    metavar="USER"
)
ap.add_argument(
    "--http-passwd",
    help="HTTP basic auth password",
    metavar="PASSWD"
)
args = ap.parse_args()
http_auth = map(lambda i: i is None, [args.http_user, args.http_passwd])
if any(http_auth) and not all(http_auth):
    ap.error("--http-user and --http-password must be both set")

opener = build_opener()
post_data = urlencode(extract_values(args.post)) if args.post is not None else args.post
headers_data = extract_values(args.headers).items() if args.headers is not None else args.headers
if headers_data is not None:
    opener.addheaders = headers_data

try:
    resp = opener.open(args.url, data=post_data)
except HTTPError as err:
    print "%d - %s" % (err.getcode(), err.reason)
    sys.exit(2)
data = resp.read()

print COLORS['green']+"----- REQUEST: -----"+COLORS['reset']
print COLORS['bgrey']+"URL: "+COLORS['bgreen']+args.url+COLORS['reset']
if post_data is not None:
    print COLORS['bgrey']+"Data: "+COLORS['bgreen']+post_data+COLORS['reset']
if headers_data is not None:
    print COLORS['bgrey']+"Headers: "+COLORS['reset']
    for k, v in headers_data:
        print (COLORS['bgreen']+"  %s: %s"+COLORS['reset']) % (k, v)
print ""
print COLORS['green']+"----- RESPONSE: -----"+COLORS['reset']
print highlight(data, JsonLexer(), TerminalFormatter())