Support AWS IMDSv2

cloud_detect/providers/aws_provider.py

@@ -1,4 +1,6 @@
+import asyncio
 import logging
+import contextlib
 from pathlib import Path
 
 import aiohttp
@@ -17,6 +19,9 @@ def __init__(self, logger=None):
         self.metadata_url = (
             'http://169.254.169.254/latest/dynamic/instance-identity/document'
         )
+        self.metadata_token_url = (
+            'http://169.254.169.254/latest/api/token'
+        )
         self.vendor_files = (
             '/sys/class/dmi/id/product_version', '/sys/class/dmi/id/bios_vendor'
         )
@@ -28,22 +33,35 @@ async def identify(self):
         self.logger.info('Try to identify AWS')
         return self.check_vendor_file() or await self.check_metadata_server()
 
-    async def check_metadata_server(self):
-        """
-            Tries to identify AWS via metadata server
-        """
-        self.logger.debug('Checking AWS metadata')
-        try:
+    async def _get_metadata_v2(self):
+        with contextlib.suppress(BaseException):
+            async with aiohttp.ClientSession() as session:
+                async with session.put(self.metadata_token_url, headers={'X-aws-ec2-metadata-token-ttl-seconds': '60'}) as response:
+                    token = await response.text()
+            return await self._get_metadata(headers={'X-aws-ec2-metadata-token': token})
+        return False
+
+    async def _get_metadata(self, headers=None):
+        with contextlib.suppress(BaseException):
             async with aiohttp.ClientSession() as session:
-                async with session.get(self.metadata_url) as response:
+                async with session.get(self.metadata_url, headers=headers) as response:
                     response = await response.json(content_type=None)
-                    if response['imageId'].startswith('ami-',) and response[
+                    if response['imageId'].startswith('ami-', ) and response[
                         'instanceId'
                     ].startswith('i-'):
                         return True
-            return False
-        except BaseException:
-            return False
+        return False
+
+    async def check_metadata_server(self):
+        """
+            Tries to identify AWS via metadata server
+        """
+        self.logger.debug('Checking AWS metadata')
+        results = await asyncio.gather(
+            self._get_metadata(),
+            self._get_metadata_v2()
+        )
+        return any(results)
 
     def check_vendor_file(self):
         """

tests/aws_provider_test.py

@@ -33,6 +33,7 @@ async def test_valid_metadata_server_check(aresponses):
 
     provider = AWSProvider()
     provider.metadata_url = f'https://{mock_host}'
+    provider.metadata_token_url = f'https://{mock_host}'
     assert await provider.check_metadata_server() is True
 
 
@@ -46,4 +47,27 @@ async def test_invalid_metadata_server_check(aresponses):
 
     provider = AWSProvider()
     provider.metadata_url = f'https://{mock_host}'
+    provider.metadata_token_url = f'https://{mock_host}'
     assert await provider.check_metadata_server() is False
+
+
+@pytest.mark.asyncio
+async def test_valid_metadata_server_check_v2(aresponses):
+    mock_host = 'testing_metadata_url.com'
+    aresponses.add(
+        mock_host, '/', 'GET',
+        aresponses.Response(status=401),
+    )
+    aresponses.add(
+        mock_host, '/', 'PUT',
+        response='token',
+    )
+    aresponses.add(
+        mock_host, '/', 'GET',
+        response={'imageId': 'ami-12312412', 'instanceId': 'i-ec12as'},
+    )
+
+    provider = AWSProvider()
+    provider.metadata_url = f'https://{mock_host}'
+    provider.metadata_token_url = f'https://{mock_host}'
+    assert await provider.check_metadata_server() is True