#! /usr/bin/env python
#
# Copyright (C) 2016 Rolf Neugebauer <rolf.neugebauer@docker.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
"""A small utility for testing if a URL works.

Returns 0 on success and 1 or 2 on failure. Returns the fetched
content on success.

Optionally, check of a the response contains a specified string.
Optionally, retry several times with a 1s sleep in between.

"""

import argparse
import ssl
import sys
import time
if sys.version_info < (3, 0):
    import urllib2 as urlreq
else:
    import urllib.request as urlreq

# Very skanky hack to disabled SSL Certificate authentication
ssl._create_default_https_context = ssl._create_unverified_context

parser = argparse.ArgumentParser()
parser.description = \
    "Fetch a URL with several retries. Optionally grep for a string"

parser.add_argument("-s", "--string",
                    help="String to search for in result")
parser.add_argument('-r', '--retry', type=int, default=0,
                    help="Number of retries")
parser.add_argument('url', nargs=1,
                    help="URL to fetch")
args = parser.parse_args()


req = urlreq.Request(args.url[0])

tries = 0
while True:
    try:
        r = urlreq.urlopen(req)
        page = str(r.read())
    except Exception as e:
        sys.stderr.write("%02d: %s %s\n" % (tries, time.ctime(), e))
        tries += 1

        if tries >= args.retry:
            sys.exit(1)
        time.sleep(1)
        continue

    print(page)
    if args.string:
        if args.string in page:
            sys.exit(0)
        else:
            sys.exit(2)
    sys.exit(0)