#!/usr/bin/python
#
# Get Cloudwatch metrics for the EBS volumes attached to an instance
#

import datetime
import logging
import sys
import urllib

import boto
from boto.exception import BotoServerError, EC2ResponseError
from boto.ec2.connection import EC2Connection


class InstanceDimension(dict):
    """
    Helper class for get_metric_statistics call
    """
    def __init__(self, name, value):
        self[name] = value


class EBSStatsCollector(object):
    """
    Retrieve EBS stats for a given instance
    """

    METRICS = [
                'VolumeWriteBytes',
                'VolumeWriteOps',
                'VolumeReadBytes',
                'VolumeReadOps',
                'VolumeQueueLength',
              ]
    """
    Amazon metrics to retrieve for each volume

    @see http://docs.amazonwebservices.com/AWSEC2/latest/UserGuide/using-cloudwatch.html
    """

    def __init__(self, minutes=60, period=60):
        """
        @param minutes: int, minutes to look back
        @param period: int, sample bucket size in seconds
        """

        self._cloudwatch = boto.connect_cloudwatch()
        self._connection = EC2Connection()
        self._minutes = minutes
        self._period = period

    def get_volume_stats(self, volume_id):
        """
        Get stats for a particular EBS volume

        @param volume_id: string, EBS volume identifier
        @return: dict, metric -> value dict
        """

        ret = {}
        for metric in self.METRICS:
            try:
                end = datetime.datetime.utcnow()
                start = end - datetime.timedelta(minutes=self._minutes)
                stats = self._cloudwatch.get_metric_statistics(self._period, start, end, metric,
                                                               'AWS/EBS', 'Average',
                                                               InstanceDimension("VolumeId", volume_id)
        )

                if not stats:
                    logging.warning('Could not get %s for volume %s (or metric is empty)', metric, volume_id)
                else:
                    ret[metric] = stats[0][u'Average']

            except BotoServerError, error:
                logging.warning('Boto API error: %s', error)
                ret[metric] = 0

        return ret

    def get_all_stats(self, instance_id):
        """
        Get CloudWatch statistics for the EBS volumes attached to a given instance


        API docs: http://docs.amazonwebservices.com/AmazonCloudWatch/latest/APIReference/API_GetMetricStatistics.html

        Hint: to figure out the exact params, use the AWS console and look
        at the query params when clicking on cloudwatch metrics...

        @param instance: string, instance to get information for
        """
        volumes = self.get_instance_volumes(instance_id)
        for (mount, volume) in volumes.items():
            print '%s (%s):' % (mount, volume.id)
            stats = self.get_volume_stats(volume.id)
            print '\t', stats

    def get_all_stats_localhost(self):
        """
        Get CloudWatch statistics for the EBS volumes attached to localhost.
        """
        return self.get_all_stats(self.get_instance_id())

    def get_instance(self, instance_id):
        """
        Return an Instance object for the given instance id

        @param instance_id: Instance id (string)
        @return: Instance object, or None if not found
        """

        try:
            reservations = self._connection.get_all_instances([instance_id])
        except EC2ResponseError, ex:
            logging.warning('Got exception when calling EC2 for instance "%s": %s',
                            instance_id, ex.error_message)
            return None

        for r in reservations:
            if len(r.instances) and r.instances[0].id == instance_id:
                return r.instances[0]

        return None

    def get_instance_volumes(self, instance_id):
        """
        Returns volumes for a specific instance as a map of mount device to volume

        @param instance_id: string, instance id
        @return: dict, mount -> volume
        """

        instance = self.get_instance(instance_id)

        devices = instance.block_device_mapping
        if not devices:
            return {}

        ebs_info = {}
        vs = self._connection.get_all_volumes()
        volumes = {}
        for volume in vs:
            volumes[volume.id] = volume
        for mount, device in devices.items():
            ebs_info[mount] = volumes[device.volume_id]

        return ebs_info

    def get_instance_id(self):
        """
        Retrieve the Amazon instance id of the instance we are running on.
        @return: string, instance id
        """
        return urllib.urlopen('http://169.254.169.254/latest/meta-data/instance-id').read()


if __name__ == "__main__":
    collector = EBSStatsCollector()
    if len(sys.argv) == 2:
        # Argument is instance id
        collector.get_all_stats(sys.argv[1])
        sys.exit(0)

    # Default to localhost
    collector.get_all_stats_localhost()