Switch to IMDSv2 and boto3

postgres-appliance/scripts/callback_aws.py

@@ -1,52 +1,41 @@
 #!/usr/bin/env python
 
-import boto.ec2
-import boto.utils
+from botocore.config import Config
+import boto3
 import logging
 import sys
-import time
+import requests
 
 logger = logging.getLogger(__name__)
 
 
-def retry(func):
-    def wrapped(*args, **kwargs):
-        count = 0
-        while True:
-            try:
-                return func(*args, **kwargs)
-            except boto.exception.BotoServerError as e:
-                if count >= 10 or str(e.error_code) not in ('Throttling', 'RequestLimitExceeded'):
-                    raise
-                logger.info('Throttling AWS API requests...')
-                time.sleep(2 ** count * 0.5)
-                count += 1
-
-    return wrapped
-
-
 def get_instance_metadata():
-    return boto.utils.get_instance_identity()['document']
+    response = requests.put(
+        url='http://169.254.169.254/latest/api/token',
+        headers={'X-aws-ec2-metadata-token-ttl-seconds': '60'}
+    )
+    token = response.text
+    instance_identity = requests.get(
+        url='http://169.254.169.254/latest/dynamic/instance-identity/document',
+        headers={'X-aws-ec2-metadata-token': token}
+    )
+    return instance_identity.json()
 
 
-@retry
 def associate_address(ec2, allocation_id, instance_id):
-    return ec2.associate_address(instance_id=instance_id, allocation_id=allocation_id, allow_reassociation=True)
+    return ec2.associate_address(InstanceId=instance_id, AllocationId=allocation_id, AllowReassociation=True)
 
 
-@retry
 def tag_resource(ec2, resource_id, tags):
-    return ec2.create_tags([resource_id], tags)
+    return ec2.create_tags(Resources=[resource_id], Tags=tags)
 
 
-@retry
 def list_volumes(ec2, instance_id):
-    return ec2.get_all_volumes(filters={'attachment.instance-id': instance_id})
+    return ec2.describe_volumes(Filters=[{'Name': 'attachment.instance-id', 'Values': [instance_id]}])
 
 
-@retry
 def get_instance(ec2, instance_id):
-    return ec2.get_only_instances([instance_id])[0]
+    return ec2.describe_instances(InstanceIds=[instance_id])['Reservations'][0]['Instances'][0]
 
 
 def main():
@@ -63,30 +52,38 @@ def main():
 
     instance_id = metadata['instanceId']
 
-    ec2 = boto.ec2.connect_to_region(metadata['region'])
+    config = Config(
+        region_name=metadata['region'],
+        retries={
+            'max_attempts': 10,
+            'mode': 'standard'
+        }
+    )
+    ec2 = boto3.client('ec2', config=config)
 
     if argc == 5 and role in ('master', 'standby_leader') and action in ('on_start', 'on_role_change'):
         associate_address(ec2, sys.argv[1], instance_id)
 
     instance = get_instance(ec2, instance_id)
 
-    tags = {'Role': role}
+    tags = [{'Key': 'Role', 'Value': role}]
     tag_resource(ec2, instance_id, tags)
 
-    tags.update({'Instance': instance_id})
+    tags.append({'Key': 'Instance', 'Value': instance_id})
 
     volumes = list_volumes(ec2, instance_id)
-    for v in volumes:
-        if 'Name' in v.tags:
+    for v in volumes['Volumes']:
+        if any(tag['Key'] == 'Name' for tag in v.get('Tags', [])):
             tags_to_update = tags
         else:
-            if v.attach_data.device == instance.root_device_name:
-                volume_device = 'root'
-            else:
-                volume_device = 'data'
-            tags_to_update = dict(tags, Name='spilo_{}_{}'.format(cluster, volume_device))
-
-        tag_resource(ec2, v.id, tags_to_update)
+            for attachment in v['Attachments']:
+                if attachment['Device'] == instance.get('RootDeviceName'):
+                    volume_device = 'root'
+                else:
+                    volume_device = 'data'
+                tags_to_update = tags + [{'Key': 'Name', 'Value': 'spilo_{}_{}'.format(cluster, volume_device)}]
+
+        tag_resource(ec2, v.get('VolumeId'), tags_to_update)
 
 
 if __name__ == '__main__':

postgres-appliance/scripts/configure_spilo.py

@@ -399,20 +399,34 @@ def get_provider():
 
     try:
         logging.info("Figuring out my environment (Google? AWS? Openstack? Local?)")
-        r = requests.get('http://169.254.169.254', timeout=2)
+        response = requests.put(
+            url='http://169.254.169.254/latest/api/token',
+            headers={'X-aws-ec2-metadata-token-ttl-seconds': '60'}
+        )
+        token = response.text
+        r = requests.get(
+            url='http://169.254.169.254',
+            headers={'X-aws-ec2-metadata-token': token}
+        )
         if r.headers.get('Metadata-Flavor', '') == 'Google':
             return PROVIDER_GOOGLE
         else:
             # accessible on Openstack, will fail on AWS
-            r = requests.get('http://169.254.169.254/openstack/latest/meta_data.json')
+            r = requests.get(
+                url='http://169.254.169.254/openstack/latest/meta_data.json',
+                headers={'X-aws-ec2-metadata-token': token}
+            )
             if r.ok:
                 # make sure the response is parsable - https://github.com/Azure/aad-pod-identity/issues/943 and
                 # https://github.com/zalando/spilo/issues/542
                 r.json()
                 return PROVIDER_OPENSTACK
 
             # is accessible from both AWS and Openstack, Possiblity of misidentification if previous try fails
-            r = requests.get('http://169.254.169.254/latest/meta-data/ami-id')
+            r = requests.get(
+                url='http://169.254.169.254/latest/meta-data/ami-id',
+                headers={'X-aws-ec2-metadata-token': token}
+            )
             return PROVIDER_AWS if r.ok else PROVIDER_UNSUPPORTED
     except (requests.exceptions.ConnectTimeout, requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout):
         logging.info("Could not connect to 169.254.169.254, assuming local Docker setup")
@@ -451,7 +465,15 @@ def get_instance_metadata(provider):
             # Try get IP via OpenStack EC2-compatible API, if can't then fail back to auto-discovered one.
             metadata['id'] = openstack_metadata['uuid']
             url = 'http://169.254.169.254/2009-04-04/meta-data'
-            r = requests.get(url)
+            response = requests.put(
+                url='http://169.254.169.254/latest/api/token',
+                headers={'X-aws-ec2-metadata-token-ttl-seconds': '60'}
+            )
+            token = response.text
+            r = requests.get(
+                url=url,
+                headers={'X-aws-ec2-metadata-token': token}
+            )
             if r.ok:
                 mapping.update({'ip': 'local-ipv4', 'id': 'instance-id'})
     else: