From aeae1d8d88699645ee9cb5e5de32c00de777cdf6 Mon Sep 17 00:00:00 2001 From: Mirela Popoveniuc Date: Wed, 24 Mar 2021 15:23:14 +0000 Subject: [PATCH 1/3] Use IMDSv2 instead of v1 when calling the EC2 Instance metadata. --- .../agent_metadata/aws_ec2_instance.py | 40 ++++++++++++++++--- .../agent_metadata/fleet_info.py | 27 +++++++++---- .../agent_metadata/test_aws_ec2_instance.py | 6 ++- 3 files changed, 59 insertions(+), 14 deletions(-) diff --git a/codeguru_profiler_agent/agent_metadata/aws_ec2_instance.py b/codeguru_profiler_agent/agent_metadata/aws_ec2_instance.py index 28fd2c7..4d15c5a 100644 --- a/codeguru_profiler_agent/agent_metadata/aws_ec2_instance.py +++ b/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. @@ -27,12 +38,29 @@ def get_fleet_instance_id(self): @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() + """ + The id of the fleet element. Eg. host name in ec2. + """ + return cls.__look_up_with_IMDSv2(EC2_HOST_NAME_URI) @classmethod def __look_up_instance_type(cls): - return http_get(url=EC2_HOST_INSTANCE_TYPE_URI).read().decode() + """ + The type of the instance. Eg. m5.2xlarge + """ + return cls.__look_up_with_IMDSv2(EC2_HOST_INSTANCE_TYPE_URI) + + @classmethod + def __look_up_with_IMDSv2(cls, url): + return http_get(url=url, + headers={EC2_METADATA_TOKEN_HEADER_KEY: cls.__look_up_ec2_api_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): @@ -45,10 +73,10 @@ def look_up_metadata(cls): 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 } diff --git a/codeguru_profiler_agent/agent_metadata/fleet_info.py b/codeguru_profiler_agent/agent_metadata/fleet_info.py index 58a5152..36756f9 100644 --- a/codeguru_profiler_agent/agent_metadata/fleet_info.py +++ b/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 } diff --git a/test/unit/agent_metadata/test_aws_ec2_instance.py b/test/unit/agent_metadata/test_aws_ec2_instance.py index 2b3d1db..a0b8ec8 100644 --- a/test/unit/agent_metadata/test_aws_ec2_instance.py +++ b/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="PARIOq_FXbIyL0maE9RcmrsyWtylvFh1ZDt0NrRUyNxeV1-DlpFpA==") yield httpretty.disable() httpretty.reset() From bdc6342e4abf873381808c23c40d5677227a77dd Mon Sep 17 00:00:00 2001 From: Mirela Popoveniuc Date: Wed, 24 Mar 2021 16:48:54 +0000 Subject: [PATCH 2/3] Replace fake token used in unit test with a dummy one. --- test/unit/agent_metadata/test_aws_ec2_instance.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/unit/agent_metadata/test_aws_ec2_instance.py b/test/unit/agent_metadata/test_aws_ec2_instance.py index a0b8ec8..82a3900 100644 --- a/test/unit/agent_metadata/test_aws_ec2_instance.py +++ b/test/unit/agent_metadata/test_aws_ec2_instance.py @@ -44,7 +44,7 @@ def around(self): httpretty.register_uri( httpretty.GET, EC2_API_TOKEN_URI, - body="PARIOq_FXbIyL0maE9RcmrsyWtylvFh1ZDt0NrRUyNxeV1-DlpFpA==") + body="fakeAndDummy") yield httpretty.disable() httpretty.reset() From bfb3206ebb753d99a22c37bc52e981b62ee73039 Mon Sep 17 00:00:00 2001 From: Mirela Popoveniuc Date: Wed, 24 Mar 2021 16:54:41 +0000 Subject: [PATCH 3/3] Get token only once. --- .../agent_metadata/aws_ec2_instance.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/codeguru_profiler_agent/agent_metadata/aws_ec2_instance.py b/codeguru_profiler_agent/agent_metadata/aws_ec2_instance.py index 4d15c5a..82b194a 100644 --- a/codeguru_profiler_agent/agent_metadata/aws_ec2_instance.py +++ b/codeguru_profiler_agent/agent_metadata/aws_ec2_instance.py @@ -37,23 +37,23 @@ def get_fleet_instance_id(self): return self.host_name @classmethod - def __look_up_host_name(cls): + 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) + return cls.__look_up_with_IMDSv2(EC2_HOST_NAME_URI, token) @classmethod - def __look_up_instance_type(cls): + 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) + return cls.__look_up_with_IMDSv2(EC2_HOST_INSTANCE_TYPE_URI, token) @classmethod - def __look_up_with_IMDSv2(cls, url): + def __look_up_with_IMDSv2(cls, url, token): return http_get(url=url, - headers={EC2_METADATA_TOKEN_HEADER_KEY: cls.__look_up_ec2_api_token()}) \ + headers={EC2_METADATA_TOKEN_HEADER_KEY: token}) \ .read().decode() @classmethod @@ -65,9 +65,10 @@ def __look_up_ec2_api_token(cls): @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 "