From 8d24b913df4572899dd11896f3344290c138c2d4 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Fri, 6 May 2022 04:22:17 -0400
Subject: [PATCH] Common: Raise exception if AWS metadata requests fail

---
 monkey/common/aws/aws_metadata.py             | 27 ++++++++-----------
 .../common/aws/test_aws_metadata.py           | 10 +++++++
 2 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/monkey/common/aws/aws_metadata.py b/monkey/common/aws/aws_metadata.py
index 2032fee81..634d41c49 100644
--- a/monkey/common/aws/aws_metadata.py
+++ b/monkey/common/aws/aws_metadata.py
@@ -14,10 +14,6 @@ logger = logging.getLogger(__name__)
 AWS_TIMEOUT = 2
 
 
-class UnknownAWSInstanceIDError(Exception):
-    """Raised if the AWS Instance ID could not be determined"""
-
-
 def fetch_aws_instance_metadata() -> Tuple[Optional[str], Optional[str], Optional[str]]:
     instance_id = None
     region = None
@@ -31,7 +27,6 @@ def fetch_aws_instance_metadata() -> Tuple[Optional[str], Optional[str], Optiona
         requests.RequestException,
         IOError,
         json.decoder.JSONDecodeError,
-        UnknownAWSInstanceIDError,
     ) as err:
         logger.debug(f"Failed init of AWSInstance while getting metadata: {err}")
         return (None, None, None)
@@ -45,20 +40,19 @@ def _fetch_aws_instance_id() -> Optional[str]:
         url,
         timeout=AWS_TIMEOUT,
     )
-
-    if not response:
-        raise UnknownAWSInstanceIDError(f"Failed fetch the AWS Instance ID from {url}")
+    response.raise_for_status()
 
     return response.text
 
 
 def _fetch_aws_region() -> Optional[str]:
-    return _parse_region(
-        requests.get(
-            AWS_LATEST_METADATA_URI_PREFIX + "meta-data/placement/availability-zone",
-            timeout=AWS_TIMEOUT,
-        ).text
+    response = requests.get(
+        AWS_LATEST_METADATA_URI_PREFIX + "meta-data/placement/availability-zone",
+        timeout=AWS_TIMEOUT,
     )
+    response.raise_for_status()
+
+    return _parse_region(response.text)
 
 
 def _parse_region(region_url_response: str) -> Optional[str]:
@@ -83,9 +77,10 @@ def _fetch_account_id() -> str:
     ../dynamic/instance-identity/document
     :return: The account id
     """
-    instance_identity_document = requests.get(
+    response = requests.get(
         AWS_LATEST_METADATA_URI_PREFIX + "dynamic/instance-identity/document",
         timeout=AWS_TIMEOUT,
-    ).text
+    )
+    response.raise_for_status()
 
-    return json.loads(instance_identity_document)[ACCOUNT_ID_KEY]
+    return json.loads(response.text)[ACCOUNT_ID_KEY]
diff --git a/monkey/tests/unit_tests/common/aws/test_aws_metadata.py b/monkey/tests/unit_tests/common/aws/test_aws_metadata.py
index 1aac906a4..e374da309 100644
--- a/monkey/tests/unit_tests/common/aws/test_aws_metadata.py
+++ b/monkey/tests/unit_tests/common/aws/test_aws_metadata.py
@@ -199,3 +199,13 @@ def test_region_not_found_request(not_found_request_mock_instance):
 
 def test_account_id_not_found_request(not_found_request_mock_instance):
     assert not_found_request_mock_instance[2] is None
+
+
+def test_instance_id_4xx_error_code():
+    with requests_mock.Mocker() as m:
+        # request made to get instance_id
+        url = f"{AWS_LATEST_METADATA_URI_PREFIX}meta-data/instance-id"
+        m.get(url, text="1234", status_code=404)
+
+        result = fetch_aws_instance_metadata()
+        assert result[0] is None