Upgrade AWS metadata to IMDSv2 (#481)

* Upgrade AWS metadata to IMDSv2

* Add minor changes to IMDS v2 upgrade PR.

Co-authored-by: Uma Annamalai <[email protected]>

Author: Robert Prast

newrelic/common/utilization.py

@@ -160,9 +160,49 @@ def detect(cls):
 class AWSUtilization(CommonUtilization):
     EXPECTED_KEYS = ('availabilityZone', 'instanceId', 'instanceType')
     METADATA_HOST = '169.254.169.254'
-    METADATA_PATH = '/2016-09-02/dynamic/instance-identity/document'
+    METADATA_PATH = '/latest/dynamic/instance-identity/document'
+    METADATA_TOKEN_PATH = '/latest/api/token'
+    HEADERS = {'X-aws-ec2-metadata-token-ttl-seconds': '21600'}
     VENDOR_NAME = 'aws'
 
+    @classmethod
+    def fetchAuthToken(cls):
+        try:
+            with cls.CLIENT_CLS(cls.METADATA_HOST,
+                                timeout=cls.FETCH_TIMEOUT) as client:
+                resp = client.send_request(method='PUT',
+                                           path=cls.METADATA_TOKEN_PATH,
+                                           params=cls.METADATA_QUERY,
+                                           headers=cls.HEADERS)
+            if not 200 <= resp[0] < 300:
+                raise ValueError(resp[0])
+            return resp[1]
+        except Exception as e:
+            _logger.debug('Unable to fetch %s data from %s%s: %r',
+                    cls.VENDOR_NAME, cls.METADATA_HOST, cls.METADATA_PATH, e)
+            return None
+
+    @classmethod
+    def fetch(cls):
+        try:
+            authToken = cls.fetchAuthToken()
+            if authToken == None:
+                return
+            cls.HEADERS = {"X-aws-ec2-metadata-token": authToken}
+            with cls.CLIENT_CLS(cls.METADATA_HOST,
+                                timeout=cls.FETCH_TIMEOUT) as client:
+                resp = client.send_request(method='GET',
+                                           path=cls.METADATA_PATH,
+                                           params=cls.METADATA_QUERY,
+                                           headers=cls.HEADERS)
+            if not 200 <= resp[0] < 300:
+                raise ValueError(resp[0])
+            return resp[1]
+        except Exception as e:
+            _logger.debug('Unable to fetch %s data from %s%s: %r',
+                    cls.VENDOR_NAME, cls.METADATA_HOST, cls.METADATA_PATH, e)
+            return None
+
 
 class AzureUtilization(CommonUtilization):
     METADATA_HOST = '169.254.169.254'