health-check.py

#!/usr/bin/env python
import argparse
import logging
import requests
import math
import os
import re
import sys
from retry.api import retry_call
from urllib.parse import urlparse

"""
	API Health Check CLI
	A small client to request AWS hosted APIs via the AWS API Gateway or direct to source.
	The client will consume exceptions during tries. If errors persist beyond tries the script
		will exit with a return code of 1 and log any exceptions to stdout.

	Example commands:

	# make AWS API Gateway health check request
	$0 --host https://host.cloud --path v1.0 --debug

	# make AWS API direct health check request
	$0 --host https://$REST_APP_ID.execute-api.eu-west-2.amazonaws.com --path path/to/feed --debug

	# make AWS API requests with tries and timeout
	# note: 5 requests over 10 seconds will equal 1 request every 2 seconds
	$0 --host [...] --path [...]  --tries 5 --timeout 10

	Requests can also be achieved using curl, when roles have been assumed
	$ curl -X GET -k -H 'x-api-key: $X_API_KEY' -i https://$REST_APP_ID.execute-api.eu-west-2.amazonaws.com/$PATH
	$ curl -X GET -k -H 'x-api-key: $X_API_KEY' -i https://host.cloud/$PATH
"""

class ApiHealthCheckException(Exception):
	""" Describes the exception raised during processing """
	def __init__(self, code, message="Command execution failed"):
		super().__init__()
		self.code = code
		self.message = message

	def __str__(self):
		return repr([self.code, str(self.message)])


class ApiHealthCheckClient:
	""" Describes a client used to check the health of AWS hosted APIs """
	API_KEY_VAR_NAME='API_HEALTH_CHECK_API_KEY'

	def __init__(self, options):
		self.client = requests
		self.logger = None
		self.config = {
			'host': None,
			'path': None,
			'auth': None,
			'debug': False
		}

		self.__configure(options)
		self.__init_logging()

	def __init_logging(self):
		if self.config.get('debug'):
			logger = logging.getLogger(__name__)
			handler = logging.StreamHandler(sys.stdout)
			handler.flush = sys.stdout.flush
			logger.addHandler(handler)
			self.logger = logger

	def log(self, message):
		""" Logs messages """
		if self.logger and self.config.get('debug'):
			self.logger.info('health-check->{}'.format(message))

	def __error(self, message):
		"""Raises ApiHealthCheckException error
			:param str message: exception message
			:raises: ApiHealthCheckException
		"""
		raise ApiHealthCheckException(1, message)

	def __sanitise_url_option(self, value):
		"""Sanitises URL options
			:param str value: option to sanitise
			:return: str -- sanitised value
		"""
		return str(value).strip('/')

	def __get_auth_config(self, key=None):
		"""Gets auth configuration in entirety or by key
			:param str key: auth object key name
			:return: mixed -- when auth key is declared return value
				(int, str, dict, etc), otherwise auth dict
		"""
		auth = self.config.get('auth')
		if key and auth:
			return auth.get(key)
		elif auth:
			return auth
		else:
			return {}

	def __get_auth_headers(self):
		headers = self.__get_auth_config('headers')
		if not headers.get('x-api-key'):
			self.log('missing API Key {}'.format(self.API_KEY_VAR_NAME))
		return headers

	def __add_host_schema(self, host):
		"""Adds desired URL schema to hosts when required
			:param str host: host value
			:return: str -- original or updated host value
		"""
		schema = re.compile('https?:\/\/.+')
		return 'https://{host}'.format_map(locals()) if not schema.match(host) else host

	def __get_url(self, config):
		"""Gets URL from config properties: host and path
			:param dict config: config properties
			:return: ParseResult -- urllib parsed URL
		"""
		host = config.get('host', '')
		path = config.get('path', '')
		if host:
			host = self.__add_host_schema(host)
			return urlparse('{host}/{path}'.format_map(locals()))

	def __configure(self, options={}):
		"""Configures a health check client based on input options
			:param dict options: HTTP method
			:raises: ApiHealthCheckException
		"""
		results = {}
		for key in self.config:
			# updates if the key exists in self.config
			# 	 and if the update value is not None
			results.update({
				k: v for k, v in options.items() if k == key and v is not None
			})

			# note: value is coerced to a string, and leading and trailing '/'
			#    are removed to clean up host and path values.
			results.update({
				'host': self.__sanitise_url_option(options.get('host')),
				'path': self.__sanitise_url_option(options.get('path'))
			})

		if results.get('host'):
			url = self.__get_url(results)

			# checks if url is supported
			if hasattr(url, 'netloc') and getattr(url, 'netloc') == '':
				self.__error('unsupported url {url}'.format_map(locals()))

			# adds the url to the results
			results.update({'url': url})

			# merges default config and results (for python >= 3.5)
			self.config = {
				**self.config,
				**results
			}
		else:
			self.__error('missing host')

	def _make_request(self, url):
		"""Makes a request against a valid URL
			:param ParseResult url: URL object
			:return: dict -- results of request
			:raises: Exception in support of various connection errors
		"""
		headers = self.__get_auth_headers()
		results = self.client.get(url.geturl(), headers=headers)
		results.raise_for_status()
		return results

	def get_resource(self):
		""" Gets responses from requests to supported APIs """
		url = self.config.get('url')

		self.log('requesting {url}'.format_map(locals()))

		try:
			# makes a request
			if hasattr(url, 'netloc'):
				return self._make_request(url)

		except Exception as error:
			self.__error(error)


class ApiHealthCheck(ApiHealthCheckClient):
	""" Describes a class used to check on AWS API health """
	def __init__(self, arguments):
		options = {
			'host': arguments.host,
			'path': arguments.path,
			'debug': arguments.debug
		}

		# adds support for an API Key when making direct requests
		# todo remove the need for an API Key with health check requests
		if os.environ.get(self.API_KEY_VAR_NAME):
			options.update({
				'auth': {
					'headers': {
						'x-api-key': os.environ[self.API_KEY_VAR_NAME]
					}
				}
			})

		# configure client
		super().__init__(options)

		# configure check
		self.config.update({
			'tries': arguments.tries,
			'timeout': arguments.timeout
		})

		if arguments.debug:
			logging.basicConfig(level=logging.INFO)

	def check(self):
		""" Checks how APIs respond to requests """
		config = self.config
		tries = config.get('tries')
		timeout = config.get('timeout')

		# setup = tries and timeout above zero
		#    with timeout devisable by tries
		setup = (tries and timeout and timeout >= tries)
		delay = math.floor(timeout / tries) if setup else None

		# note: tries defaults to 1, overrides zero tries with default
		tries = tries if delay else 1
		delay = delay if delay else 1

		self.log('using tries {tries} timeout {timeout} delay {delay}'.format_map(locals()))

		# gets the resource with or without
		#    tries and timeout configuration
		results = retry_call(
			self.get_resource,
			tries=tries,
			delay=delay,
			exceptions=ApiHealthCheckException
		)

		# show results
		print(repr(results))


def parse_args():
	""" Processes command line arguments """
	parser = argparse.ArgumentParser(description='Checks the health of a specified endpoint')
	parser.add_argument('--host', help='The host - can also include URL schema, domain and port')
	parser.add_argument('--path', help='The path (URI)')
	parser.add_argument('--tries', help='Connect tries before failing the check', type=int, default=1)
	parser.add_argument('--timeout', help='Timeout (seconds) before failing the check', type=int, default=1)
	parser.add_argument('--debug', help='Enable debug', action='store_true')
	return parser.parse_args()

if __name__ == '__main__':
	args = parse_args()
	healthCheck = ApiHealthCheck(args)
	healthCheck.check()