smoke.py

#!/usr/bin/env python

# Sample YAML file that smoke.py expects:
#
# example.com:
#   - alice
#   - bob/marley
#   - charlie
# example.org:
#   - david
#   - ""
#
# and would produce requests to example.com/alice, example.com/bob/marley,
# example.com/charlie, example.org/david, and example.org/

import yaml
import argparse
import requests


def run(endpoints, protocol="http"):
    """Iterates over YAML dictionary of lists to pass to smoke()
    Args:
        endpoints (str): name of file containing YAML dictionary
        protocol (Optional[str]): name of HTTP protocol to check with
    """

    # read from this file
    with open(endpoints, "r") as f:
        servers = yaml.load(f)

    # iterate over each server list
    for server in servers:
        # iterate over each endpoint in server
        for uri in servers[server]:
            smoke(protocol=protocol,
                  server=server,
                  uri=uri)


def smoke(protocol, server, uri):
    """Makes request to server with URI over protocol (on localhost).
    Prints result to stdout. If protocol is "all", recurses on each
    protocol.
    Args:
        protocol (str): name of HTTP protocol or "all" for both
        server (str): the server name that is listening for requests
        uri (str): path
    """

    if protocol == "all":
        for proto in ["http", "https"]:

            # recurse to check each protocol
            smoke(protocol=proto,
                  server=server,
                  uri=uri)
    else:
        # get return status code
        code = requests.get("{}://localhost/{}".format(protocol, uri),
                            headers={'Host': server}).status_code

        # print results
        print("{} {}://{}/{}".format(code, protocol, server, uri))


def main():

    parser = argparse.ArgumentParser()

    parser.add_argument("endpoints",
                        help="YAML file of servers with lists of endpoints")
    parser.add_argument("-p", "--proto",
                        choices=["all", "http", "https"],
                        default="http",
                        help="protocol(s) to use when testing endpoints")

    args = parser.parse_args()

    run(endpoints=args.endpoints, protocol=args.proto)


if __name__ == "__main__":
    main()