From 8cebcfb19500dac080c73ff1013d566c20a37671 Mon Sep 17 00:00:00 2001
From: troymc <troy@ascribe.io>
Date: Tue, 29 Mar 2016 11:53:38 +0200
Subject: [PATCH] Acquire & assign elastic IPs the boto3 way

---
 deploy-cluster-aws/get_elastic_ips.py | 76 +++++++++++++--------------
 deploy-cluster-aws/startup.sh         |  3 +-
 2 files changed, 39 insertions(+), 40 deletions(-)

diff --git a/deploy-cluster-aws/get_elastic_ips.py b/deploy-cluster-aws/get_elastic_ips.py
index 73daf9aa..9976eb27 100644
--- a/deploy-cluster-aws/get_elastic_ips.py
+++ b/deploy-cluster-aws/get_elastic_ips.py
@@ -1,16 +1,17 @@
 # -*- coding: utf-8 -*-
 
-# create Elastic IPs and assign them to instances if needed
+""" Create Elastic IPs and assign them to instances if needed.
+"""
 
 from __future__ import unicode_literals
-import json
 import os
-import boto.ec2
+import boto3
 import argparse
 import time
 
-AWS_ACCESS_KEY_ID = os.environ["AWS_ACCESS_KEY_ID"]
-AWS_SECRET_ACCESS_KEY = os.environ["AWS_SECRET_ACCESS_KEY"]
+AWS_ACCESS_KEY_ID = os.environ['AWS_ACCESS_KEY_ID']
+AWS_SECRET_ACCESS_KEY = os.environ['AWS_SECRET_ACCESS_KEY']
+AWS_REGION = os.environ['AWS_REGION']
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--tag", help="tag instances in aws")
@@ -20,46 +21,45 @@ if args.tag:
     tag = args.tag
 else:
     # reading credentials from config for remote connection
-    print('usage: python3 get_elastic_ips.py --tag <tag>')
+    print('usage: python get_elastic_ips.py --tag <tag>')
     print('reason: tag missing!!!')
     exit(1)
 
-conn = boto.ec2.connect_to_region("eu-central-1",
-            aws_access_key_id=AWS_ACCESS_KEY_ID,
-            aws_secret_access_key=AWS_SECRET_ACCESS_KEY)
-
-INSTANCE_IDS = []
+# Connect to Amazon EC2
+ec2 = boto3.resource(service_name='ec2',
+                     region_name=AWS_REGION,
+                     aws_access_key_id=AWS_ACCESS_KEY_ID,
+                     aws_secret_access_key=AWS_SECRET_ACCESS_KEY)
 
 
-def prepare_list(tag):
-    reservations = conn.get_all_instances(filters={"tag:Name" : tag})
-    instances = [i for r in reservations for i in r.instances]
-    for i in instances:
-        inst = i.__dict__
-        #print (inst)
-        #break
-        inst_id = inst.get('id')
+# Get a list of all instances with the specified tag.
+# (Technically, instances_with_tag is an ec2.instancesCollection.)
+instances_with_tag = ec2.instances.filter(
+    Filters=[{'Name': 'tag:Name', 'Values': [tag]}]
+    )
 
-        INSTANCE_IDS.append(inst_id)
-    return INSTANCE_IDS
+print('Allocating elastic IP addresses and assigning them to the instances...')
 
+for instance in instances_with_tag:
+    # Create a client from the ec2 resource
+    # See http://boto3.readthedocs.org/en/latest/guide/clients.html
+    client = ec2.meta.client
 
-def get_new_pubDNS():
-    eip = conn.allocate_address()
-    return eip
+    # Acquire an Elastic IP address
+    # response is a dict. See http://tinyurl.com/z2n7u9k
+    response = client.allocate_address(DryRun=False, Domain='standard')
+    public_ip = response['PublicIp']
+    print('public_ip = {}'.format(public_ip))
 
-if __name__ == "__main__":
-    # hostlist.tmp (JSON) erzeugen
-    instlist = prepare_list(tag)
+    # Associate that Elastic IP address with an instance
+    response2 = client.associate_address(
+        DryRun=False,
+        InstanceId=instance.instance_id,
+        PublicIp=public_ip
+        )
+    print('was associated with the instance with id {}'.
+          format(instance.instance_id))
 
-    for entry in range(0,len(instlist)):
-
-        instance_id = instlist[entry]
-        print(instance_id)
-        newpubDNS = get_new_pubDNS()
-        inID = str(newpubDNS).split(':')[1]
-        print(inID)
-        conn.associate_address(instance_id, public_ip=inID)
-
-    # make sure all addresse are assigned...
-    time.sleep(30)
+# Make sure all IP addresses are assigned...
+print('Waiting 30 seconds to make sure all IP addresses are assigned...')
+time.sleep(30)
diff --git a/deploy-cluster-aws/startup.sh b/deploy-cluster-aws/startup.sh
index c97a6b33..059b576d 100755
--- a/deploy-cluster-aws/startup.sh
+++ b/deploy-cluster-aws/startup.sh
@@ -62,9 +62,8 @@ python run_and_tag.py --tag $TAG --nodes $NODES
 python wait_until_all_running.py --tag $TAG
 
 # in case of elastic ips...
-python3 get_elastic_ips.py --tag $TAG
+python get_elastic_ips.py --tag $TAG
 
-# everything prepared. now wait until instances up and running!
 # generate hostlist.py and add_keys.sh
 python3 create_hostlist.py --tag $TAG > hostlist.py
 # make add_keys executable and execute