s3-canary.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
# Loic Lambiel ©
# License MIT

import argparse
import logging
import logging.handlers
import socket
import sys
import time

try:
    import boto
    import boto.s3.connection
    from boto.s3.key import Key
except ImportError:
    print("It look like boto module isn't installed. \
          Please install it using pip install boto")
    sys.exit(1)


try:
    import bernhard
except ImportError:
    print("It look like riemann client (bernard) isn't installed. \
          Please install it using pip install bernhard")
    sys.exit(1)


try:
    from configparser import ConfigParser
except ImportError:  # python 2
    from ConfigParser import ConfigParser


class CustomError(Exception):

    def __init__(self, arg):
        self.msg = arg


def main():
    parser = argparse.ArgumentParser(description='''
This script create a file, read it and delete it on a s3 compliant
storage. If any error occur during the process, an alarm is being
sent to riemann monitoring. time metric is also sent to riemann''')
    parser.add_argument('-version',
                        action='version',
                        version='%(prog)s 1.0, Loic Lambiel, exoscale')
    parser.add_argument('-key',
                        help='S3 user key',
                        required=True,
                        type=str,
                        dest='key')
    parser.add_argument('-secret',
                        help='S3 user secret',
                        required=True,
                        type=str,
                        dest='secret')
    parser.add_argument('-host',
                        help='S3 host',
                        required=True,
                        type=str,
                        dest='host')
    parser.add_argument('-bucket',
                        help='S3 bucket',
                        required=True,
                        type=str,
                        dest='bucket')
    parser.add_argument('-env',
                        help='Environnement used by the riemann service',
                        required=True,
                        type=str,
                        dest='env')
    parser.add_argument('-alertstate',
                        help='Alert level to raise if the test fails',
                        required=False,
                        type=str,
                        default='critical',
                        dest='state')
    args = vars(parser.parse_args())
    return args


def s3test(args):
    KEY = args['key']
    SECRET = args['secret']
    HOST = args['host']
    BUCKET = args['bucket']
    ENV = args['env']

    logging.basicConfig(format=('%(asctime)s %(pathname)s '
                                '%(levelname)s:%(message)s'),
                        level=logging.DEBUG,
                        filename=logfile)
    logging.getLogger().addHandler(logging.StreamHandler())
    logging.getLogger('boto').setLevel(logging.DEBUG)
    logging.info('Test started for env %s', ENV)

    conn = boto.connect_s3(
        aws_access_key_id=KEY,
        aws_secret_access_key=SECRET,
        host=HOST,
        calling_format=boto.s3.connection.OrdinaryCallingFormat(),
    )

    logging.info('Creating bucket %s', BUCKET)
    bucket = conn.create_bucket(BUCKET)
    logging.info('Bucket create for env %s', ENV)

    k = Key(bucket)
    k.key = 's3-canary'
    logging.info('Writing file')
    k.set_contents_from_string('This is a test of S3')
    logging.info('Write done for env %s', ENV)
    time.sleep(3)
    logging.info('Reading back the file for env %s', ENV)

    if k.get_contents_as_string() != "This is a test of S3":
        raise CustomError('Failed to read back the created file !')

    logging.info('File has been read back for env %s', ENV)
    logging.info('Deleting bucket for env %s', ENV)
    bucket.delete_key(k)
    logging.info('Bucket deleted for env %s', ENV)


# main
if __name__ == "__main__":
    args = main()
    region = args['host'].split('.')[0]
    logfile = "/var/log/s3-canary-{}.log".format(region)
    conf = ConfigParser()
    conf.read(("/etc/bernhard.conf",))

    client = bernhard.SSLClient(host=conf.get('default', 'riemann_server'),
                                port=int(conf.get('default', 'riemann_port')),
                                keyfile=conf.get('default', 'tls_cert_key'),
                                certfile=conf.get('default', 'tls_cert'),
                                ca_certs=conf.get('default', 'tls_ca_cert'))
    ENV = args['env']
    HOST = args['host']
    state = args['state']
    exectimeservice = "{}.s3_canary.exectime".format(HOST)
    checkservice = "{}.s3_canary.check".format(HOST)
    start_time = time.time()
    try:
        s3test(args)
        exectime = time.time() - start_time
        host = socket.gethostname()
        client.send({'host': host,
                     'service': exectimeservice,
                     'state': 'ok',
                     'tags': ['duration', ENV],
                     'ttl': 600,
                     'metric': exectime})
        client.send({'host': host,
                     'service': checkservice,
                     'state': 'ok',
                     'tags': ['s3_canary.py', 'duration', ENV],
                     'ttl': 600,
                     'metric': 0})

        logging.info('Script completed successfully for env %s', ENV)

    except Exception as e:
        logging.exception("An exception occured. Exception is: %s", e)
        host = socket.gethostname()
        exectime = 61
        txt = "An exception occurred on s3_canary.py: {}. \
               See logfile {} for more info".format(e, logfile)
        client.send({'host': host,
                     'service': checkservice,
                     'description': txt,
                     'state': state,
                     'tags': ['s3_canary.py', 'duration', ENV],
                     'ttl': 600,
                     'metric': 1})
        client.send({'host': host,
                     'service': exectimeservice,
                     'state': 'ok',
                     'tags': ['duration', ENV],
                     'ttl': 600,
                     'metric': exectime})
        raise