Use IMDSv2

IMDSv1 has problems and nowadays IMDSv2 is preferred.

Signed-off-by: Pablo Barbáchano <pablob@amazon.com>

Author: pb8o

tests/framework/utils_cpuid.py

@@ -7,6 +7,7 @@
 from enum import Enum, auto
 
 from framework.utils import run_cmd
+from framework.utils_imdsv2 import imdsv2_get
 import host_tools.network as net_tools
 
 ARM_CPU_DICT = {"0xd0c": "ARM_NEOVERSE_N1"}
@@ -46,10 +47,8 @@ def get_cpu_model_name():
 
 
 def get_instance_type():
-    """Get the instance type through IMDS."""
-    imds_cmd = "curl http://169.254.169.254/latest/meta-data/instance-type"
-    _, stdout, _ = run_cmd(imds_cmd)
-    return stdout
+    """Get the instance type through IMDSv2"""
+    return imdsv2_get("/meta-data/instance-type")
 
 
 def check_guest_cpuid_output(

tests/framework/utils_imdsv2.py

@@ -0,0 +1,63 @@
+# Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""A simple IMDSv2 client
+
+- https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
+- https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-categories.html
+
+Important! For this client to work in a container scenario, make sure your
+instances are set with an adequate hop limit (2 for example). See
+`ec2:MetadataHttpPutResponseHopLimit`
+"""
+
+
+import time
+
+import requests
+
+IMDSV2_HDR_TOKEN_TTL = "X-aws-ec2-metadata-token-ttl-seconds"
+IMDSV2_HDR_TOKEN = "X-aws-ec2-metadata-token"
+
+
+class IMDSv2Client:
+    """
+    A simple IMDSv2 client.
+
+    >>> IMDSv2Client().get("/meta-data/instance-type")
+    """
+
+    def __init__(self, endpoint="http://169.254.169.254", version="latest"):
+        self.endpoint = endpoint
+        self.version = version
+        self.ttl = 21600
+        self.token_expiry_time = 0
+        self.token = None
+
+    def get_token(self):
+        """Get a token from IMDSv2"""
+        if self.token_expiry_time < time.time():
+            headers = {IMDSV2_HDR_TOKEN_TTL: str(self.ttl)}
+            # To get a token, docs say to always use latest
+            url = f"{self.endpoint}/latest/api/token"
+            res = requests.put(url, headers=headers)
+            self.token = res.content
+            self.token_expiry_time = time.time() + self.ttl
+        return self.token
+
+    def get(self, path):
+        """
+        Get a metadata path from IMDSv2
+
+        >>> IMDSv2Client().get("/meta-data/instance-type")
+        """
+        headers = {IMDSV2_HDR_TOKEN: self.get_token()}
+        url = f"{self.endpoint}/{self.version}{path}"
+        res = requests.get(url, headers=headers)
+        if res.status_code != 200:
+            raise Exception(f"IMDSv2 returned {res.status_code} for {url}")
+        return res.text
+
+
+IMDS_V2 = IMDSv2Client()
+imdsv2_get = IMDS_V2.get