From f9fa9381efb7fb8e37baba2aa4c36cf5348f143b Mon Sep 17 00:00:00 2001 From: "Gareth J. Greenaway" Date: Fri, 3 Nov 2023 13:41:56 -0700 Subject: [PATCH] Account for situation where the metadata grain fails because the AWS environment requires an authentication token to query the metadata URL. --- salt/grains/metadata.py | 55 ++++++- tests/pytests/unit/grains/test_metadata.py | 180 +++++++++++++++++++++ 2 files changed, 231 insertions(+), 4 deletions(-) create mode 100644 tests/pytests/unit/grains/test_metadata.py diff --git a/salt/grains/metadata.py b/salt/grains/metadata.py index d1dcd41643b..ba5eddaf9c0 100644 --- a/salt/grains/metadata.py +++ b/salt/grains/metadata.py @@ -24,7 +24,7 @@ import salt.utils.stringutils # metadata server information IP = "169.254.169.254" -HOST = "http://{}/".format(IP) +HOST = f"http://{IP}/" def __virtual__(): @@ -36,16 +36,55 @@ def __virtual__(): if result != 0: return False if http.query(os.path.join(HOST, "latest/"), status=True).get("status") != 200: - return False + # Initial connection failed, might need a token + _refresh_token() + if ( + http.query( + os.path.join(HOST, "latest/"), + status=True, + header_dict={ + "X-aws-ec2-metadata-token": __context__["metadata_aws_token"] + }, + ).get("status") + != 200 + ): + return False return True +def _refresh_token(): + __context__["metadata_aws_token"] = http.query( + os.path.join(HOST, "latest/api/token"), + method="PUT", + header_dict={"X-aws-ec2-metadata-token-ttl-seconds": "21600"}, + ).get("body") + + def _search(prefix="latest/"): """ Recursively look up all grains in the metadata server """ ret = {} - linedata = http.query(os.path.join(HOST, prefix), headers=True) + if "metadata_aws_token" in __context__: + if ( + http.query( + os.path.join(HOST, "latest/"), + status=True, + header_dict={ + "X-aws-ec2-metadata-token": __context__["metadata_aws_token"] + }, + ).get("status") + != 200 + ): + _refresh_token() + + linedata = http.query( + os.path.join(HOST, prefix), + header_dict={"X-aws-ec2-metadata-token": __context__["metadata_aws_token"]}, + headers=True, + ) + else: + linedata = http.query(os.path.join(HOST, prefix), headers=True) if "body" not in linedata: return ret body = salt.utils.stringutils.to_unicode(linedata["body"]) @@ -68,7 +107,15 @@ def _search(prefix="latest/"): key, value = line.split("=") ret[value] = _search(prefix=os.path.join(prefix, key)) else: - retdata = http.query(os.path.join(HOST, prefix, line)).get("body", None) + if "metadata_aws_token" in __context__: + retdata = http.query( + os.path.join(HOST, prefix, line), + header_dict={ + "X-aws-ec2-metadata-token": __context__["metadata_aws_token"] + }, + ).get("body", None) + else: + retdata = http.query(os.path.join(HOST, prefix, line)).get("body", None) # (gtmanfred) This try except block is slightly faster than # checking if the string starts with a curly brace if isinstance(retdata, bytes): diff --git a/tests/pytests/unit/grains/test_metadata.py b/tests/pytests/unit/grains/test_metadata.py new file mode 100644 index 00000000000..ddce7d14dba --- /dev/null +++ b/tests/pytests/unit/grains/test_metadata.py @@ -0,0 +1,180 @@ +""" + Unit test for salt.grains.metadata + + + :codeauthor: :email" `Gareth J. Greenaway + +""" + +import logging + +import pytest + +import salt.grains.metadata as metadata +import salt.utils.http as http +from tests.support.mock import MagicMock, create_autospec, patch + +# from Exception import Exception, ValueError + +log = logging.getLogger(__name__) + + +class MockSocketClass: + def __init__(self, *args, **kwargs): + pass + + def settimeout(self, *args, **kwargs): + pass + + def connect_ex(self, *args, **kwargs): + return 0 + + +@pytest.fixture +def configure_loader_modules(): + return {metadata: {"__opts__": {"metadata_server_grains": "True"}}} + + +def test_metadata_search(): + def mock_http( + url="", + method="GET", + headers=False, + header_list=None, + header_dict=None, + status=False, + ): + metadata_vals = { + "http://169.254.169.254/latest/api/token": { + "body": "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX==", + "status": 200, + "headers": {}, + }, + "http://169.254.169.254/latest/": { + "body": "meta-data", + "headers": {}, + }, + "http://169.254.169.254/latest/meta-data/": { + "body": "ami-id\nami-launch-index\nami-manifest-path\nhostname", + "headers": {}, + }, + "http://169.254.169.254/latest/meta-data/ami-id": { + "body": "ami-xxxxxxxxxxxxxxxxx", + "headers": {}, + }, + "http://169.254.169.254/latest/meta-data/ami-launch-index": { + "body": "0", + "headers": {}, + }, + "http://169.254.169.254/latest/meta-data/ami-manifest-path": { + "body": "(unknown)", + "headers": {}, + }, + "http://169.254.169.254/latest/meta-data/hostname": { + "body": "ip-xx-x-xx-xx.us-west-2.compute.internal", + "headers": {}, + }, + } + + return metadata_vals[url] + + with patch( + "salt.utils.http.query", + create_autospec(http.query, autospec=True, side_effect=mock_http), + ): + ret = metadata.metadata() + assert ret == { + "meta-data": { + "ami-id": "ami-xxxxxxxxxxxxxxxxx", + "ami-launch-index": "0", + "ami-manifest-path": "(unknown)", + "hostname": "ip-xx-x-xx-xx.us-west-2.compute.internal", + } + } + + with patch.dict( + metadata.__context__, + { + "metadata_aws_token": "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX==" + }, + ): + with patch( + "salt.utils.http.query", + create_autospec(http.query, autospec=True, side_effect=mock_http), + ): + ret = metadata.metadata() + assert ret == { + "meta-data": { + "ami-id": "ami-xxxxxxxxxxxxxxxxx", + "ami-launch-index": "0", + "ami-manifest-path": "(unknown)", + "hostname": "ip-xx-x-xx-xx.us-west-2.compute.internal", + } + } + + +def test_metadata_refresh_token(): + with patch( + "salt.utils.http.query", + create_autospec( + http.query, + autospec=True, + return_value={ + "body": "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX==", + }, + ), + ): + metadata._refresh_token() + assert "metadata_aws_token" in metadata.__context__ + assert ( + metadata.__context__["metadata_aws_token"] + == "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX==" + ) + + +def test_metadata_virtual(): + with patch("socket.socket", MagicMock(return_value=MockSocketClass())): + with patch( + "salt.utils.http.query", + create_autospec( + http.query, + autospec=True, + return_value={"error": "[Errno -2] Name or service not known"}, + ), + ): + assert metadata.__virtual__() is False + + with patch( + "salt.utils.http.query", + create_autospec( + http.query, + autospec=True, + return_value={ + "body": "dynamic\nmeta-data\nuser-data", + "status": 200, + }, + ), + ): + assert metadata.__virtual__() is True + + with patch( + "salt.utils.http.query", + create_autospec( + http.query, + autospec=True, + side_effect=[ + { + "body": "", + "status": 401, + }, + { + "body": "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX==", + }, + { + "body": "dynamic\nmeta-data\nuser-data", + "status": 200, + }, + ], + ), + ): + assert metadata.__virtual__() is True