~linaro-aws-devs/linaro-aws-tools/trunk

#!/usr/bin/python

from boto import connect_ec2, connect_s3
from boto.ec2 import regions
from datetime import datetime
import unittest

# objects which exist for OLD_RUNNING_DAYS or more are considered old
OLD_RUNNING_DAYS = 5

# all elastic addresses across regions
EC2_ADDRESSES = set((
    # region, ip, instance

    # android-build.linaro.org
    ("us-east-1", "50.17.250.69", "i-874e5fe9"),

    # ci.linaro.org
    ("us-east-1", "50.17.200.206", "i-4b36752f"),

    # Bip
    ("us-east-1", "50.19.223.142", "i-3935a658"),

    # cards.linaro.org
    ("us-east-1", "107.21.205.76", "i-89eb05ed"),

    # flexlm.linaro.org
    ("us-east-1", "23.21.178.170", "i-f3bd4997"),
    ))

EC2_INSTANCES_DATA = (
    # zone, instance, root device, type, key, status (optional), tags

    # android-build.linaro.org
    {'placement': 'us-east-1b', 'id': 'i-874e5fe9',
     'root_device_type': 'ebs',
     'instance_type': 'm1.small', 'key_name': 'mwh@grond', 'state': 'running',
     'tags': {'Name': 'android-build.linaro.org'}},

    # ci.linaro.org
    {'placement': 'us-east-1b', 'id': 'i-4b36752f',
     'root_device_type': 'ebs',
     'instance_type': 'm1.medium', 'key_name': 'pfalcon@home', 'state': 'running',
     'tags': {'Name': 'ci.linaro.org'}},

    # bip.linaro.org
    {'placement': 'us-east-1b', 'id': 'i-3935a658',
     'root_device_type': 'ebs',
     'instance_type': 't1.micro', 'key_name': 'lool@bee',
     'tags': {'Name': 'bip.linaro.org'}},

    # michaelh's personal micro-instance
    {'placement': 'us-east-1b', 'id': 'i-6aa9240a',
     'root_device_type': 'ebs',
     'instance_type': 't1.micro', 'key_name': 'michaelh@crucis',
     'tags': {'Name': "michaelh's personal micro-instance"}},

    # cards.linaro.org
    {'placement': 'us-east-1b', 'id': 'i-89eb05ed',
     'root_device_type': 'ebs',
     'instance_type': 'c1.medium', 'key_name': 'lool@bee',
     'tags': {'Name': 'cards.linaro.org'}},

    # flexlm.linaro.org
    {'placement': 'us-east-1b', 'id': 'i-f3bd4997',
     'root_device_type': 'ebs',
     'instance_type': 't1.micro', 'key_name': 'lool@bee',
     'tags': {'Name': 'flexlm.linaro.org'}},

    # FastModels hackbox
    {'placement': 'us-east-1b', 'id': 'i-95ac58f1',
     'root_device_type': 'ebs',
     'instance_type': 'm2.xlarge', 'key_name': 'lool@bee',
     'tags': {'Name': 'FastModels hackbox'}},

    # linaro-ci-dashboard demo
    {'placement': 'us-east-1b', 'id': 'i-d6f791ac',
     'root_device_type': 'ebs',
     'instance_type': 't1.micro', 'key_name': 'pfalcon@home',
     'tags': {'Name': 'linaro-ci-dashboard demo'}},
    )

# XXX no support for egress filtering
EC2_SECURITY_GROUPS_DATA = (
    # name, description, region, rules
    # rules are a list of IP permissions which include IP protocol, a port
    # range for TCP and UDP or type and code for ICMP, and a source grant in
    # the form of a source CIDR or of a group-name followed by an AWS id

    # XXX used by default notably for jenkins slaves; these should use a
    # dedicated sg but jenkins lacks this feature
    {'region': 'us-east-1', 'name': 'default',
     'description': 'default group',
     'rules': (
        # allow incoming TCP from other instances in the default group
        {'ip_protocol': 'tcp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        # allow incoming UDP from other instances in the default group
        {'ip_protocol': 'udp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        # allow incoming ICMP from other instances in the default group
        {'ip_protocol': 'icmp', 'port_range': '-1--1',
         'grant': 'default-228690440174'},
        # incoming SSH from jenkins-master group
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': 'jenkins-master-228690440174'},
        # incoming SSH from anywhere
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
     ),
    },

    # default groups in some other regions; not all have our customized rule
    # for SSH access
    {'region': 'ap-northeast-1', 'name': 'default',
     'description': 'default group',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'udp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'icmp', 'port_range': '-1--1',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'ap-southeast-1', 'name': 'default',
     'description': 'default group',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'udp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'icmp', 'port_range': '-1--1',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'eu-west-1', 'name': 'default',
     'description': 'default group',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'udp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'icmp', 'port_range': '-1--1',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'sa-east-1', 'name': 'default',
     'description': 'default group',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'udp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'icmp', 'port_range': '-1--1',
         'grant': 'default-228690440174'},
     ),
    },
    {'region': 'us-west-1', 'name': 'default',
     'description': 'default group',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'udp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'icmp', 'port_range': '-1--1',
         'grant': 'default-228690440174'},
     ),
    },
    {'region': 'us-west-2', 'name': 'default',
     'description': 'default group',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'udp', 'port_range': '0-65535',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'icmp', 'port_range': '-1--1',
         'grant': 'default-228690440174'},
     ),
    },

    # per service-type
    {'region': 'us-east-1', 'name': 'bip',
     'description': 'Bip',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
        {'ip_protocol': 'tcp', 'port_range': '6667-6667',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'us-east-1', 'name': 'flexlm-server',
     'description': 'FlexLM licensing server',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '8224-8224',
         'grant': 'flexlm-slave-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '18224-18224',
         'grant': 'flexlm-slave-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'us-east-1', 'name': 'flexlm-slave',
     'description': 'FlexLM licensing slave',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'us-east-1', 'name': 'git-mirror',
     'description': 'Git mirror',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '3128-3128',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '3129-3129',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '9418-9418',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '8080-8080',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'us-east-1', 'name': 'jenkins-master',
     'description': 'Jenkins master',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '3128-3128',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
        {'ip_protocol': 'tcp', 'port_range': '80-80',
         'grant': '0.0.0.0/0'},
        {'ip_protocol': 'tcp', 'port_range': '443-443',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'us-east-1', 'name': 'patch-tracker',
     'description': 'Patch Tracker',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '80-80',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'us-east-1', 'name': 't-git-mirror',
     'description': 'Cloud Buildd git mirror',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': 'default-228690440174'},
        {'ip_protocol': 'tcp', 'port_range': '22-22',
         'grant': '0.0.0.0/0'},
     ),
    },
    {'region': 'us-east-1', 'name': 'web-server',
     'description': 'Web server',
     'rules': (
        {'ip_protocol': 'tcp', 'port_range': '80-80',
         'grant': '0.0.0.0/0'},
     ),
    },

)

# TODO: add zone; EBS volumes need to be in the same zone as the instance
# using them anyway
EC2_VOLUMES = set((
    # region, volume, instance, size

    # android-build.linaro.org
    ("us-east-1", "vol-e9270382", "i-874e5fe9", 8),
    ("us-east-1", "vol-1ee7ad71", "i-874e5fe9", 2),
    ("us-east-1", "vol-9a612ff5", "i-874e5fe9", 500),

    # ci.linaro.org
    ("us-east-1", "vol-c787a9ab", "i-4b36752f", 8),
    ("us-east-1", "vol-949db0fb", "i-4b36752f", 2),
    ("us-east-1", "vol-58240237", "i-4b36752f", 300),

    # Bip
    ("us-east-1", "vol-39153e52", "i-3935a658", 8),

    # michaelh's personal micro-instance
    ("us-east-1", "vol-f44a399e", "i-6aa9240a", 8),

    # cards.linaro.org
    ("us-east-1", "vol-23a09b4e", "i-89eb05ed", 8),

    # flexlm.linaro.org
    ("us-east-1", "vol-3fac3253", "i-f3bd4997", 8),

    # FastModels hackbox
    ("us-east-1", "vol-4fa63823", "i-95ac58f1", 8),

    # linaro-ci-dashboard demo
    ("us-east-1", "vol-6a91c311", "i-d6f791ac", 8),
    ))

S3_BUCKETS = set((
    # name

    # seed shared repository tarball to optimize bzr checkouts of lp:gcc-linaro
    "gcc-linaro",

    # toolchain outputs
    'linaro-toolchain-files',
    'linaro-toolchain-sources',
    'linaro-toolchain-builds',

    # billing reports
    'linaro-billing-reports',
    ))

# caches
_ALL_EC2_ADDRESSES = None
_ALL_EC2_INSTANCES = None
_ALL_EC2_VOLUMES = None
_ALL_EC2_SECURITY_GROUPS = None
# TODO: EC2 reserved instances
# TODO: EC2 snapshots
_ALL_S3_BUCKETS = None

def _gather_data():
    global _ALL_EC2_ADDRESSES, _ALL_EC2_INSTANCES
    global _ALL_EC2_SECURITY_GROUPS, _ALL_EC2_VOLUMES
    global _ALL_S3_BUCKETS

    addresses = []
    instances = []
    security_groups = []
    volumes = []
    for region in regions():
        ec2 = connect_ec2(region=region)

        addresses.extend(ec2.get_all_addresses())
        security_groups.extend(ec2.get_all_security_groups())
        volumes.extend(ec2.get_all_volumes())
        for reservation in ec2.get_all_instances():
            instances.extend(reservation.instances)

    s3 = connect_s3()
    buckets = s3.get_all_buckets()

    _ALL_EC2_ADDRESSES = addresses
    _ALL_EC2_INSTANCES = instances
    _ALL_EC2_SECURITY_GROUPS = security_groups
    _ALL_EC2_VOLUMES = volumes
    _ALL_S3_BUCKETS = buckets

def get_all_addresses():
    global _ALL_EC2_ADDRESSES

    if _ALL_EC2_ADDRESSES is None:
        _gather_data()
    return _ALL_EC2_ADDRESSES

def get_all_instances():
    global _ALL_EC2_INSTANCES

    if _ALL_EC2_INSTANCES is None:
        _gather_data()
    return _ALL_EC2_INSTANCES

def get_instances():
    instances = []
    for instance in get_all_instances():
        if instance.state != 'terminated':
            instances.append(instance)
    return instances

def get_all_security_groups():
    global _ALL_EC2_SECURITY_GROUPS

    if _ALL_EC2_SECURITY_GROUPS is None:
        _gather_data()
    return _ALL_EC2_SECURITY_GROUPS

def get_all_volumes():
    global _ALL_EC2_VOLUMES

    if _ALL_EC2_VOLUMES is None:
        _gather_data()
    return _ALL_EC2_VOLUMES

def get_all_buckets():
    global _ALL_S3_BUCKETS

    if _ALL_S3_BUCKETS is None:
        _gather_data()
    return _ALL_S3_BUCKETS

def formated_table(column_sizes, data):
    print " ".join(
        [heading.ljust(size)[:size] for heading, size in column_sizes])
    print
    for line in data:
        assert len(line) == len(column_sizes)
        columns = []
        for i in range(len(column_sizes)):
            size = column_sizes[i][1]
            columns.append(str(line[i]).ljust(size)[:size])
        print " ".join(columns)

def show_addresses():
    column_sizes = (("region", 9), ("ip", 15), ("instance", 10))
    data = [
        (a.region.name, a.public_ip, a.instance_id)
            for a in get_all_addresses()]
    formated_table(column_sizes, data)

def show_instances():
    column_sizes = (
        ("zone", 10), ("instance", 10), ("type", 9), ("status", 7),
        ("key", 17), ("launch time", 19))
    data = [
        (i.placement, i.id, i.instance_type, i.state, i.key_name,
         i.launch_time)
            for i in get_all_instances()]
    formated_table(column_sizes, data)

def show_security_groups():
    column_sizes = (
        ("region", 9), ("name", 14), ("proto", 5), ("ports", 7),
        ("grant", 27))
    data = [
        (sg.region.name, sg.name, r.ip_protocol,
         (r.from_port == r.to_port) and r.from_port
             or "%s-%s" % (r.from_port, r.to_port),
         g)
            for sg in get_all_security_groups()
                for r in sg.rules
                    for g in r.grants]
    formated_table(column_sizes, data)

def show_volumes():
    column_sizes = (
        ("region", 9), ("volume", 12), ("instance", 10) , ("size", 4),
        ("creation time", 19))
    data = [
        (v.region.name, v.id, v.attach_data.instance_id, str(v.size),
         v.create_time)
            for v in get_all_volumes()]
    formated_table(column_sizes, data)

def show_buckets():
    column_sizes = (("name", 30), ("creation date", 19))
    data = [(b.name, b.creation_date) for b in get_all_buckets()]
    formated_table(column_sizes, data)

def get_xsd_datetime(xsd_datetime):
    # strip ".000Z"
    return datetime.strptime(xsd_datetime[:-5], "%Y-%m-%dT%H:%M:%S")

def is_old(xsd_datetime):
    delta = datetime.now() - get_xsd_datetime(xsd_datetime)
    return delta.days >= OLD_RUNNING_DAYS

def get_old_instances():
    instances = []
    for instance in get_instances():
        if is_old(instance.launch_time):
            instances.append(instance)
    return instances

def get_old_volumes():
    volumes = []
    for volume in get_all_volumes():
        if is_old(volume.create_time):
            volumes.append(volume)
    return volumes

def get_old_buckets():
    buckets = []
    for bucket in get_all_buckets():
        if is_old(bucket.creation_date):
            buckets.append(bucket)
    return buckets

def instance_matches_data(instance, data):
    for key in data.keys():
        inst_key = key
        # boto 1.9b used rootDeviceType while 2.0b4 uses root_device_type
        if key == "root_device_type" and not hasattr(instance, key):
            inst_key = "rootDeviceType"
        if not hasattr(instance, inst_key):
            return False
        if getattr(instance, inst_key) != data[key]:
            return False
    return True

class TestEC2Addresses(unittest.TestCase):
    def test_elastic_addresses_present(self):
        all_addresses = get_all_addresses()
        s = set([
            (a.region.name, a.public_ip, a.instance_id)
                for a in all_addresses])
        self.assertEqual(EC2_ADDRESSES, s)

class TestEC2Instances(unittest.TestCase):
    def test_instances_present(self):
        all_instances = get_instances()
        no_match = []
        multiple_matches = []
        for instance_data in EC2_INSTANCES_DATA:
            matching = 0
            for i in all_instances:
                if instance_matches_data(i, instance_data):
                    matching += 1
            if matching == 0:
                no_match.append(instance_data)
            elif matching > 1:
                multiple_matches.append(instance_data)
        self.assertEqual([], no_match)
        self.assertEqual([], multiple_matches)

    # TODO this duplicates code from test_instances_present()
    def test_no_unknown_old_instance(self):
        no_match = []
        multiple_matches = []
        for i in get_old_instances():
            matching = 0
            for instance_data in EC2_INSTANCES_DATA:
                if instance_matches_data(i, instance_data):
                    matching += 1
            if matching == 0:
                no_match.append(i)
            elif matching > 1:
                multiple_matches.append(i)
        self.assertEqual([], no_match)
        self.assertEqual([], multiple_matches)

class TestEC2SecurityGroups(unittest.TestCase):
    def test_security_groups_present(self):
        all_security_groups = get_all_security_groups()
        s = set([tuple(map(str,
            (sg.region.name, sg.name, r.ip_protocol,
             '%s-%s' % (r.from_port, r.to_port), str(g), sg.description)))
                for sg in all_security_groups
                    for r in sg.rules
                        for g in r.grants])
        expected = set([
            (sg['region'], sg['name'], r['ip_protocol'], r['port_range'],
             r['grant'], sg['description'])
                for sg in EC2_SECURITY_GROUPS_DATA
                    for r in sg['rules']])
        self.assertEqual(expected, s)

class TestEC2Volumes(unittest.TestCase):
    def get_volumes_set(self, volumes):
        return set([
            (v.region.name, v.id, v.attach_data.instance_id, v.size)
                for v in volumes])

    def test_no_unknown_old_volume(self):
        old_volumes_set = self.get_volumes_set(get_old_volumes())
        self.assertEqual(set(), old_volumes_set.difference(EC2_VOLUMES))

    def test_volumes_present(self):
        all_volumes_set = self.get_volumes_set(get_all_volumes())
        self.assertEqual(set(), EC2_VOLUMES.difference(all_volumes_set))

class TestS3Buckets(unittest.TestCase):
    def get_buckets_set(self, buckets):
        return set([b.name for b in buckets])

    def test_no_unknown_old_bucket(self):
        old_buckets_set = self.get_buckets_set(get_old_buckets())
        self.assertEqual(set(), old_buckets_set.difference(S3_BUCKETS))

    def test_buckets_present(self):
        all_buckets_set = self.get_buckets_set(get_all_buckets())
        self.assertEqual(set(), S3_BUCKETS.difference(all_buckets_set))

if __name__ == "__main__":
    show_addresses()
    print
    print
    show_instances()
    print
    print
    show_security_groups()
    print
    print
    show_volumes()
    print
    print
    show_buckets()
    print
    print
    unittest.main()