Merge pull request #24 from aws/mirelap-IMDSv2-PR

Author: mirelap-amazon

codeguru_profiler_agent/agent_metadata/aws_ec2_instance.py

@@ -9,8 +9,19 @@
 EC2_HOST_NAME_URI = DEFAULT_EC2_METADATA_URI + "local-hostname"
 EC2_HOST_INSTANCE_TYPE_URI = DEFAULT_EC2_METADATA_URI + "instance-type"
 
+# Used for IMDSv2 to retrieve API token that will be used to call the EC2 METADATA service.
+# https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
+# Bandit marks the following line as risky because it contains the word "token",
+# thought it doesn't contain any secret; ignoring with # nosec
+# https://bandit.readthedocs.io/en/latest/plugins/b105_hardcoded_password_string.html
+EC2_API_TOKEN_URI = "http://169.254.169.254/latest/api/token"  # nosec
+EC2_METADATA_TOKEN_HEADER_KEY = 'X-aws-ec2-metadata-token'  # nosec
+EC2_METADATA_TOKEN_TTL_HEADER_KEY = 'X-aws-ec2-metadata-token-ttl-seconds'  # nosec
+EC2_METADATA_TOKEN_TTL_HEADER_VALUE = '21600'  # nosec
+
 logger = logging.getLogger(__name__)
 
+
 class AWSEC2Instance(FleetInfo):
     """
     This class will get and parse the EC2 metadata if available.
@@ -26,29 +37,47 @@ def get_fleet_instance_id(self):
         return self.host_name
 
     @classmethod
-    def __look_up_host_name(cls):
-        # The id of the fleet element. Eg. host name in ec2.
-        return http_get(url=EC2_HOST_NAME_URI).read().decode()
+    def __look_up_host_name(cls, token):
+        """
+        The id of the fleet element. Eg. host name in ec2.
+        """
+        return cls.__look_up_with_IMDSv2(EC2_HOST_NAME_URI, token)
+
+    @classmethod
+    def __look_up_instance_type(cls, token):
+        """
+        The type of the instance. Eg. m5.2xlarge
+        """
+        return cls.__look_up_with_IMDSv2(EC2_HOST_INSTANCE_TYPE_URI, token)
 
     @classmethod
-    def __look_up_instance_type(cls):
-        return http_get(url=EC2_HOST_INSTANCE_TYPE_URI).read().decode()
+    def __look_up_with_IMDSv2(cls, url, token):
+        return http_get(url=url,
+                        headers={EC2_METADATA_TOKEN_HEADER_KEY: token}) \
+            .read().decode()
+
+    @classmethod
+    def __look_up_ec2_api_token(cls):
+        return http_get(url=EC2_API_TOKEN_URI,
+                        headers={EC2_METADATA_TOKEN_TTL_HEADER_KEY: EC2_METADATA_TOKEN_TTL_HEADER_VALUE}) \
+            .read().decode()
 
     @classmethod
     def look_up_metadata(cls):
         try:
+            token = cls.__look_up_ec2_api_token()
             return cls(
-                host_name=cls.__look_up_host_name(),
-                host_type=cls.__look_up_instance_type()
+                host_name=cls.__look_up_host_name(token),
+                host_type=cls.__look_up_instance_type(token)
             )
         except Exception:
             log_exception(logger, "Unable to get Ec2 instance metadata, this is normal when running in a different "
                                   "environment (e.g. Fargate), profiler will still work")
             return None
-    
+
     def serialize_to_map(self):
         return {
-            "computeType": "aws_ec2_instance", 
-            "hostName": self.host_name, 
+            "computeType": "aws_ec2_instance",
+            "hostName": self.host_name,
             "hostType": self.host_type
         }

codeguru_profiler_agent/agent_metadata/fleet_info.py

@@ -1,17 +1,29 @@
-from abc import ABCMeta, abstractmethod
+import logging
+import os
 import uuid
+
+from abc import ABCMeta, abstractmethod
 from urllib import request
-import os
+from urllib.request import Request
 
 METADATA_URI_TIMEOUT_SECONDS = 3
 
-def http_get(url):
+logger = logging.getLogger(__name__)
+
+
+def http_get(url, headers={}):
     # request.urlopen has been flagged as risky because it can be used to open local files if url starts with
     # file://, to protect us from that we add a check in the passed url.
     # With this check we can tell bandit (static analysis tool) to ignore this error with #nosec
+    # https://bandit.readthedocs.io/en/latest/blacklists/blacklist_calls.html#b310-urllib-urlopen
+
+    logger.debug("Making a request to {} with headers set for these keys: {}".format(url, headers.keys()))
     if not url.startswith("http"):
         raise ValueError("url for metadata is not a valid http address. We will not try to get metadata")
-    return request.urlopen(url, timeout=METADATA_URI_TIMEOUT_SECONDS)  # nosec
+    req = Request(url)
+    for key in headers:
+        req.add_header(key, headers[key])
+    return request.urlopen(req, timeout=METADATA_URI_TIMEOUT_SECONDS)  # nosec
 
 
 class FleetInfo(metaclass=ABCMeta):  # pragma: no cover
@@ -23,6 +35,7 @@ def __init__(self):
     This id can be the hostname for an EC2 or the task ARN for Fargate.
     @return the id in string.
     """
+
     @abstractmethod
     def get_fleet_instance_id(self):
         pass
@@ -42,7 +55,7 @@ class DefaultFleetInfo(FleetInfo):
 
     def __init__(self):
         self.fleet_instance_id = str(uuid.uuid4())
-        try: 
+        try:
             # sched_getaffinity gives the number of logical cpus that the process can use on Unix systems.
             # If not available, we default to the cpu_count().
             self.vCPUs = len(os.sched_getaffinity(0))
@@ -54,7 +67,7 @@ def get_fleet_instance_id(self):
 
     def serialize_to_map(self):
         return {
-            "id": self.fleet_instance_id, 
-            "type": "UNKNOWN", 
+            "id": self.fleet_instance_id,
+            "type": "UNKNOWN",
             "vCPUs": self.vCPUs
         }

test/unit/agent_metadata/test_aws_ec2_instance.py

@@ -2,7 +2,7 @@
 import httpretty
 import sys
 from codeguru_profiler_agent.agent_metadata.aws_ec2_instance import EC2_HOST_NAME_URI, \
-    EC2_HOST_INSTANCE_TYPE_URI
+    EC2_HOST_INSTANCE_TYPE_URI, EC2_API_TOKEN_URI
 from codeguru_profiler_agent.agent_metadata.aws_ec2_instance import AWSEC2Instance
 
 
@@ -41,6 +41,10 @@ def around(self):
                     httpretty.GET,
                     EC2_HOST_INSTANCE_TYPE_URI,
                     body="testHostType")
+                httpretty.register_uri(
+                    httpretty.GET,
+                    EC2_API_TOKEN_URI,
+                    body="fakeAndDummy")
                 yield
                 httpretty.disable()
                 httpretty.reset()