diff --git a/setup.py b/setup.py index 3de7591..09c0330 100755 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ """ import os -from setuptools import setup +from setuptools import setup, find_packages __author__ = 'Siemens AG' @@ -56,8 +56,12 @@ def read(fname): "aws", "snapshot", "bucket", "ebs", "s3" ], python_requires=">=3.5", - package_dir={"": "src"}, - packages=["snap_to_bucket"], + package_dir={ + "snap_to_bucket": "src/snap_to_bucket", + "snap_to_bucket.handlers": "src/snap_to_bucket/handlers", + "snap_to_bucket.runner": "src/snap_to_bucket/runner" + }, + packages=find_packages("./src"), install_requires=[ 'boto3', 'psutil', diff --git a/src/snap_to_bucket/handlers/ec_2_handler.py b/src/snap_to_bucket/handlers/ec_2_handler.py index 63a8f06..3969a2d 100644 --- a/src/snap_to_bucket/handlers/ec_2_handler.py +++ b/src/snap_to_bucket/handlers/ec_2_handler.py @@ -10,6 +10,7 @@ import os import sys +import http import json import math import urllib.request @@ -17,6 +18,9 @@ import boto3 +AWS_META_URL = "http://169.254.169.254/latest/dynamic/instance-identity/document/" +AWS_TOKEN_URL = "http://169.254.169.254/latest/api/token" + class Ec2Handler: """ @@ -57,10 +61,14 @@ def __init__(self, tag="snap-to-bucket", volume_type="gp2", verbose=0, self.iops = iops self.throughput = throughput try: - response = urllib.request.urlopen( - "http://169.254.169.254/latest/dynamic/instance-identity/document/") + response = urllib.request.urlopen(AWS_META_URL) + except urllib.error.HTTPError as ex: + if ex.code == http.client.UNAUTHORIZED: + response = self.get_info_from_imds2() + else: + raise ex except urllib.error.URLError as ex: - print("Script needs to run on an EC2 instance", file=sys.stderr) + print("Script needs to run on an EC2 instance!", file=sys.stderr) raise ex self.instance_info = json.loads(response.read().decode("UTF-8").strip()) response.close() @@ -70,6 +78,33 @@ def __init__(self, tag="snap-to-bucket", volume_type="gp2", verbose=0, os.environ["AWS_DEFAULT_REGION"] = self.instance_info["region"] self.ec2client = boto3.client("ec2") + def get_info_from_imds2(self): + """ + Try to get the instance from IMDSv2. + + The function gets token from AWS and use it to authenticate call to + IMDSv2. + + :return: Response from AWS meta data url + :raises urllib.error.URLError: If call fails + """ + try: + if self.verbose > 0: + print("Unable to get instance info. Trying with IMDSv2.") + req = urllib.request.Request(AWS_TOKEN_URL, + headers={"X-aws-ec2-metadata-token-ttl-seconds": 60}, + method="PUT") + response = urllib.request.urlopen(req) + token = response.read().decode("UTF-8").strip() + response.close() + req = urllib.request.Request(AWS_META_URL, + headers={"X-aws-ec2-metadata-token": token}) + toreturn = urllib.request.urlopen(req) + except urllib.error.URLError as ex: + print("Script needs to run on an EC2 instance!", file=sys.stderr) + raise ex + return toreturn + def get_snapshots(self): """ Get the list of snapshots which are to be migrated.