chore: validate iam_host (#471)

Author: karenc-bq

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -125,7 +125,7 @@ IamAuthPlugin.InvalidPort=[IamAuthPlugin] Port number: {} is not valid. Port num
 IamAuthPlugin.NoValidPort=[IamAuthPlugin] Unable to determine a valid port.
 IamAuthPlugin.UnhandledException=[IamAuthPlugin] Unhandled exception: {}
 IamAuthPlugin.UseCachedIamToken=[IamAuthPlugin] Used cached IAM token = {}
-
+IAMAuthPlugin.InvalidHost=[IamAuthPlugin] Invalid IAM host {}. The IAM host must be a valid RDS or Aurora endpoint.
 IamPlugin.IsNoneOrEmpty=[IamPlugin] Property "{}" is None or empty.
 
 LogUtils.Topology=[LogUtils] Topology {}

aws_advanced_python_wrapper/utils/iamutils.py

@@ -17,6 +17,11 @@
 from datetime import datetime
 from typing import TYPE_CHECKING
 
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.utils.messages import Messages
+from aws_advanced_python_wrapper.utils.rds_url_type import RdsUrlType
+from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.hostinfo import HostInfo
 
@@ -25,10 +30,21 @@
 
 
 class IamAuthUtils:
-
     @staticmethod
     def get_iam_host(props: Properties, host_info: HostInfo):
-        return WrapperProperties.IAM_HOST.get(props) if WrapperProperties.IAM_HOST.get(props) else host_info.host
+        host = WrapperProperties.IAM_HOST.get(props) if WrapperProperties.IAM_HOST.get(props) else host_info.host
+        IamAuthUtils.validate_iam_host(host)
+        return host
+
+    @staticmethod
+    def validate_iam_host(host: str | None):
+        if host is None:
+            raise AwsWrapperError(Messages.get_formatted("IAMAuthPlugin.InvalidHost", "[No host provided]"))
+
+        utils = RdsUtils()
+        rds_type = utils.identify_rds_type(host)
+        if rds_type == RdsUrlType.OTHER or rds_type == RdsUrlType.IP_ADDRESS:
+            raise AwsWrapperError(Messages.get_formatted("IAMAuthPlugin.InvalidHost", host))
 
     @staticmethod
     def get_port(props: Properties, host_info: HostInfo, dialect_default_port: int) -> int:

tests/integration/container/test_iam_authentication.py

@@ -73,8 +73,7 @@ def test_iam_wrong_database_username(self, test_environment: TestEnvironment,
                 plugins="iam",
                 **props)
 
-    def test_iam_no_database_username(self, test_environment: TestEnvironment,
-                                      test_driver: TestDriver, conn_utils, props):
+    def test_iam_no_database_username(self, test_driver: TestDriver, conn_utils, props):
         target_driver_connect = DriverHelper.get_connect_func(test_driver)
         params = conn_utils.get_connect_params()
         params.pop("use_pure", None)  # AWS tokens are truncated when using the pure Python MySQL driver
@@ -83,6 +82,15 @@ def test_iam_no_database_username(self, test_environment: TestEnvironment,
         with pytest.raises(AwsWrapperError):
             AwsWrapperConnection.connect(target_driver_connect, **params, plugins="iam", **props)
 
+    def test_iam_invalid_host(self, test_driver: TestDriver, conn_utils, props):
+        target_driver_connect = DriverHelper.get_connect_func(test_driver)
+        params = conn_utils.get_connect_params()
+        params.pop("use_pure", None)  # AWS tokens are truncated when using the pure Python MySQL driver
+        params.update({"iam_host": "<>", "plugins": "iam"})
+
+        with pytest.raises(AwsWrapperError):
+            AwsWrapperConnection.connect(target_driver_connect, **params, **props)
+
     def test_iam_using_ip_address(self, test_environment: TestEnvironment,
                                   test_driver: TestDriver, conn_utils, props):
         target_driver_connect = DriverHelper.get_connect_func(test_driver)

tests/unit/test_iam_plugin.py

@@ -21,6 +21,7 @@
 
 import pytest
 
+from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.hostinfo import HostInfo
 from aws_advanced_python_wrapper.iam_plugin import IamAuthPlugin, TokenInfo
 from aws_advanced_python_wrapper.utils.properties import (Properties,
@@ -350,10 +351,17 @@ def test_connect_with_specified_region(mocker, mock_plugin_service, mock_session
     mock_dialect.set_password.assert_called_with(expected_props, f"{_TEST_TOKEN}:{iam_region}")
 
 
+@pytest.mark.parametrize("iam_host", [
+    pytest.param("foo.testdb.us-east-2.rds.amazonaws.com"),
+    pytest.param("test.cluster-123456789012.us-east-2.rds.amazonaws.com"),
+    pytest.param("test-.cluster-ro-123456789012.us-east-2.rds.amazonaws.com"),
+    pytest.param("test.cluster-custom-123456789012.us-east-2.rds.amazonaws.com"),
+    pytest.param("test-.proxy-123456789012.us-east-2.rds.amazonaws.com.cn"),
+    pytest.param("test-.proxy-123456789012.us-east-2.rds.amazonaws.com.proxy"),
+])
 @patch("aws_advanced_python_wrapper.iam_plugin.IamAuthPlugin._token_cache", _token_cache)
-def test_connect_with_specified_host(mocker, mock_plugin_service, mock_session, mock_func, mock_client, mock_dialect):
+def test_connect_with_specified_host(iam_host: str, mocker, mock_plugin_service, mock_session, mock_func, mock_client, mock_dialect):
     test_props: Properties = Properties({"user": "postgresqlUser"})
-    iam_host: str = "foo.testdb.us-east-2.rds.amazonaws.com"
 
     test_props[WrapperProperties.IAM_HOST.name] = iam_host
 
@@ -365,7 +373,7 @@ def test_connect_with_specified_host(mocker, mock_plugin_service, mock_session,
     target_plugin.connect(
         target_driver_func=mocker.MagicMock(),
         driver_dialect=mock_dialect,
-        host_info=HostInfo("pg.testdb.us-east-2.rds.amazonaws.com"),
+        host_info=HostInfo("bar.foo.com"),
         props=test_props,
         is_initial_connection=False,
         connect_func=mock_func)
@@ -376,16 +384,38 @@ def test_connect_with_specified_host(mocker, mock_plugin_service, mock_session,
         DBUsername="postgresqlUser"
     )
 
-    actual_token = _token_cache.get("us-east-2:foo.testdb.us-east-2.rds.amazonaws.com:5432:postgresqlUser")
+    actual_token = _token_cache.get(f"us-east-2:{iam_host}:5432:postgresqlUser")
+    assert actual_token is not None
     assert _GENERATED_TOKEN != actual_token.token
-    assert f"{_TEST_TOKEN}:foo.testdb.us-east-2.rds.amazonaws.com" == actual_token.token
+    assert f"{_TEST_TOKEN}:{iam_host}" == actual_token.token
     assert actual_token.is_expired() is False
 
-    # Assert password has been updated to the value in token cache
-    expected_props = {"iam_host": "foo.testdb.us-east-2.rds.amazonaws.com", "user": "postgresqlUser"}
-    mock_dialect.set_password.assert_called_with(expected_props, f"{_TEST_TOKEN}:foo.testdb.us-east-2.rds.amazonaws.com")
-
 
 def test_aws_supported_regions_url_exists():
     url = "https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/Concepts.RegionsAndAvailabilityZones.html"
     assert 200 == urllib.request.urlopen(url).getcode()
+
+
+@pytest.mark.parametrize("host", [
+    pytest.param("<>"),
+    pytest.param("#"),
+    pytest.param("'"),
+    pytest.param("\""),
+    pytest.param("%"),
+    pytest.param("^"),
+    pytest.param("https://foo.com/abc.html"),
+    pytest.param("foo.boo//"),
+    pytest.param("8.8.8.8"),
+    pytest.param("a.b"),
+])
+def test_invalid_iam_host(host, mocker, mock_plugin_service, mock_session, mock_func, mock_client, mock_dialect):
+    test_props: Properties = Properties({"user": "postgresqlUser"})
+    with pytest.raises(AwsWrapperError):
+        target_plugin: IamAuthPlugin = IamAuthPlugin(mock_plugin_service, mock_session)
+        target_plugin.connect(
+            target_driver_func=mocker.MagicMock(),
+            driver_dialect=mock_dialect,
+            host_info=HostInfo(host),
+            props=test_props,
+            is_initial_connection=False,
+            connect_func=mock_func)