test(auth): IAM credential authentication

Author: Brooke-white

test/integration/plugin/test_credentials_providers.py

@@ -113,16 +113,15 @@ def testSslAndIam(idp_arg):
     idp_arg["iam"] = True
     with pytest.raises(
         redshift_connector.InterfaceError,
-        match="Invalid connection property setting. SSL must be enabled when using IAM",
+        match="Invalid connection property setting",
     ):
         redshift_connector.connect(**idp_arg)
 
     idp_arg["iam"] = False
     idp_arg["credentials_provider"] = "OktacredentialSProvider"
     with pytest.raises(
         redshift_connector.InterfaceError,
-        match="Invalid connection property setting. IAM must be enabled when using credentials "
-        "via identity provider",
+        match="Invalid connection property setting",
     ):
         redshift_connector.connect(**idp_arg)
 
@@ -131,7 +130,7 @@ def testSslAndIam(idp_arg):
     idp_arg["credentials_provider"] = None
     with pytest.raises(
         redshift_connector.InterfaceError,
-        match="Invalid connection property setting. " "Credentials provider cannot be None when IAM is enabled",
+        match="Invalid connection property setting",
     ):
         redshift_connector.connect(**idp_arg)
 

test/unit/auth/__init__.py

@@ -0,0 +1,123 @@
+import typing
+from unittest.mock import MagicMock, patch
+
+import boto3  # type: ignore
+import pytest  # type: ignore
+from pytest_mock import mocker
+
+from redshift_connector import InterfaceError
+from redshift_connector.auth.aws_credentials_provider import (
+    AWSCredentialsProvider,
+    AWSDirectCredentialsHolder,
+)
+from redshift_connector.credentials_holder import AWSProfileCredentialsHolder
+from redshift_connector.redshift_property import RedshiftProperty
+
+
+def _make_aws_credentials_obj_with_profile() -> AWSCredentialsProvider:
+    cred_provider: AWSCredentialsProvider = AWSCredentialsProvider()
+    rp: RedshiftProperty = RedshiftProperty()
+    profile_name: str = "myProfile"
+
+    rp.profile = profile_name
+
+    cred_provider.add_parameter(rp)
+    return cred_provider
+
+
+def _make_aws_credentials_obj_with_credentials() -> AWSCredentialsProvider:
+    cred_provider: AWSCredentialsProvider = AWSCredentialsProvider()
+    rp: RedshiftProperty = RedshiftProperty()
+    access_key_id: str = "my_access"
+    secret_key: str = "my_secret"
+    session_token: str = "my_session"
+
+    rp.access_key_id = access_key_id
+    rp.secret_access_key = secret_key
+    rp.session_token = session_token
+
+    cred_provider.add_parameter(rp)
+    return cred_provider
+
+
+def test_create_aws_credentials_provider_with_profile():
+    cred_provider: AWSCredentialsProvider = _make_aws_credentials_obj_with_profile()
+    assert cred_provider.profile == "myProfile"
+    assert cred_provider.access_key_id is None
+    assert cred_provider.secret_access_key is None
+    assert cred_provider.session_token is None
+
+
+def test_create_aws_credentials_provider_with_credentials():
+    cred_provider: AWSCredentialsProvider = _make_aws_credentials_obj_with_credentials()
+    assert cred_provider.profile is None
+    assert cred_provider.access_key_id == "my_access"
+    assert cred_provider.secret_access_key == "my_secret"
+    assert cred_provider.session_token == "my_session"
+
+
+def test_get_cache_key_with_profile():
+    cred_provider: AWSCredentialsProvider = _make_aws_credentials_obj_with_profile()
+    assert cred_provider.get_cache_key() == hash(cred_provider.profile)
+
+
+def test_get_cache_key_with_credentials():
+    cred_provider: AWSCredentialsProvider = _make_aws_credentials_obj_with_credentials()
+    assert cred_provider.get_cache_key() == hash("my_access")
+
+
+def test_get_credentials_checks_cache_first(mocker):
+    mocked_credential_holder = MagicMock()
+
+    def mock_set_cache(cp: AWSCredentialsProvider, key: str = "tomato") -> None:
+        cp.cache[key] = mocked_credential_holder  # type: ignore
+
+    cred_provider: AWSCredentialsProvider = AWSCredentialsProvider()
+    mocker.patch("redshift_connector.auth.AWSCredentialsProvider.get_cache_key", return_value="tomato")
+
+    with patch("redshift_connector.auth.AWSCredentialsProvider.refresh") as mocked_refresh:
+        mocked_refresh.side_effect = mock_set_cache(cred_provider)
+        get_cache_key_spy = mocker.spy(cred_provider, "get_cache_key")
+
+        assert cred_provider.get_credentials() == mocked_credential_holder
+
+        assert get_cache_key_spy.called is True
+        assert get_cache_key_spy.call_count == 1
+
+
+def test_get_credentials_refresh_error_is_raised(mocker):
+    cred_provider: AWSCredentialsProvider = AWSCredentialsProvider()
+    mocker.patch("redshift_connector.auth.AWSCredentialsProvider.get_cache_key", return_value="tomato")
+    expected_exception = "something went wrong"
+
+    with patch("redshift_connector.auth.AWSCredentialsProvider.refresh") as mocked_refresh:
+        mocked_refresh.side_effect = Exception(expected_exception)
+
+        with pytest.raises(InterfaceError, match=expected_exception):
+            cred_provider.get_credentials()
+
+
+def test_refresh_uses_profile_if_present(mocker):
+    cred_provider: AWSCredentialsProvider = _make_aws_credentials_obj_with_profile()
+    mocked_boto_session: MagicMock = MagicMock()
+
+    with patch("boto3.Session", return_value=mocked_boto_session):
+        cred_provider.refresh()
+
+        assert hash("myProfile") in cred_provider.cache
+        assert isinstance(cred_provider.cache[hash("myProfile")], AWSProfileCredentialsHolder)
+        assert cred_provider.cache[hash("myProfile")].profile == "myProfile"
+
+
+def test_refresh_uses_credentials_if_present(mocker):
+    cred_provider: AWSCredentialsProvider = _make_aws_credentials_obj_with_credentials()
+    mocked_boto_session: MagicMock = MagicMock()
+
+    with patch("boto3.Session", return_value=mocked_boto_session):
+        cred_provider.refresh()
+
+        assert hash("my_access") in cred_provider.cache
+        assert isinstance(cred_provider.cache[hash("my_access")], AWSDirectCredentialsHolder)
+        assert cred_provider.cache[hash("my_access")].access_key_id == "my_access"
+        assert cred_provider.cache[hash("my_access")].secret_access_key == "my_secret"
+        assert cred_provider.cache[hash("my_access")].session_token == "my_session"

test/unit/auth/test_aws_credentials_provider.py

@@ -38,8 +38,7 @@ def test_ssl_and_iam_invalid_should_fail(idp_arg):
     idp_arg["credentials_provider"] = "OktacredentialSProvider"
     with pytest.raises(
         redshift_connector.InterfaceError,
-        match="Invalid connection property setting. IAM must be enabled when using credentials "
-        "via identity provider",
+        match="Invalid connection property setting",
     ):
         redshift_connector.connect(**idp_arg)
 
@@ -48,6 +47,6 @@ def test_ssl_and_iam_invalid_should_fail(idp_arg):
     idp_arg["credentials_provider"] = None
     with pytest.raises(
         redshift_connector.InterfaceError,
-        match="Invalid connection property setting. " "Credentials provider cannot be None when IAM is enabled",
+        match="Invalid connection property setting",
     ):
         redshift_connector.connect(**idp_arg)

test/unit/plugin/test_credentials_providers.py

@@ -0,0 +1,65 @@
+import typing
+from unittest.mock import MagicMock
+
+import pytest  # type: ignore
+from pytest_mock import mocker
+
+from redshift_connector.credentials_holder import (
+    ABCAWSCredentialsHolder,
+    AWSDirectCredentialsHolder,
+    AWSProfileCredentialsHolder,
+)
+
+
+def test_aws_direct_credentials_holder_should_have_session():
+    mocked_session: MagicMock = MagicMock()
+    obj: AWSDirectCredentialsHolder = AWSDirectCredentialsHolder(
+        access_key_id="something", secret_access_key="secret", session_token="fornow", session=mocked_session
+    )
+
+    assert isinstance(obj, ABCAWSCredentialsHolder)
+    assert hasattr(obj, "get_boto_session")
+    assert obj.has_associated_session == True
+    assert obj.get_boto_session() == mocked_session
+
+
+valid_aws_direct_credential_params: typing.List[typing.Dict[str, typing.Optional[str]]] = [
+    {"access_key_id": "something", "secret_access_key": "secret", "session_token": "fornow"},
+    {"access_key_id": "something", "secret_access_key": "secret", "session_token": None},
+]
+
+
+@pytest.mark.parametrize("input", valid_aws_direct_credential_params)
+def test_aws_direct_credentials_holder_get_session_credentials(input):
+    input["session"] = MagicMock()
+    obj: AWSDirectCredentialsHolder = AWSDirectCredentialsHolder(**input)
+
+    ret_value: typing.Dict[str, str] = obj.get_session_credentials()
+
+    assert len(ret_value) == 3 if input["session_token"] is not None else 2
+
+    assert ret_value["aws_access_key_id"] == input["access_key_id"]
+    assert ret_value["aws_secret_access_key"] == input["secret_access_key"]
+
+    if input["session_token"] is not None:
+        assert ret_value["aws_session_token"] == input["session_token"]
+
+
+def test_aws_profile_credentials_holder_should_have_session():
+    mocked_session: MagicMock = MagicMock()
+    obj: AWSProfileCredentialsHolder = AWSProfileCredentialsHolder(profile="myprofile", session=mocked_session)
+
+    assert isinstance(obj, ABCAWSCredentialsHolder)
+    assert hasattr(obj, "get_boto_session")
+    assert obj.has_associated_session == True
+    assert obj.get_boto_session() == mocked_session
+
+
+def test_aws_profile_credentials_holder_get_session_credentials():
+    profile_val: str = "myprofile"
+    obj: AWSProfileCredentialsHolder = AWSProfileCredentialsHolder(profile=profile_val, session=MagicMock())
+
+    ret_value = obj.get_session_credentials()
+    assert len(ret_value) == 1
+
+    assert ret_value["profile"] == profile_val

test/unit/test_credentials_holder.py

@@ -4,6 +4,7 @@
 from pytest_mock import mocker
 
 from redshift_connector import InterfaceError, RedshiftProperty, set_iam_properties
+from redshift_connector.auth import AWSCredentialsProvider
 from redshift_connector.config import ClientProtocolVersion
 from redshift_connector.iam_helper import set_iam_credentials
 from redshift_connector.plugin import (
@@ -40,7 +41,7 @@ def mock_all_provider_get_credentials(mocker):
         mocker.patch("redshift_connector.plugin.{}.get_credentials".format(provider), return_value=None)
 
 
-def get_set_iam_properties_args(**kwargs):
+def get_set_iam_properties_args(**kwargs) -> typing.Dict[str, typing.Any]:
     return {
         "info": RedshiftProperty(),
         "user": "awsuser",
@@ -80,6 +81,10 @@ def get_set_iam_properties_args(**kwargs):
         "allow_db_user_override": True,
         "client_protocol_version": ClientProtocolVersion.BASE_SERVER,
         "database_metadata_current_db_only": True,
+        "access_key_id": None,
+        "secret_access_key": None,
+        "session_token": None,
+        "profile": None,
         "ssl_insecure": None,
         **kwargs,
     }
@@ -138,19 +143,75 @@ def test_set_iam_properties_enforce_client_protocol_version(_input):
     ({"ssl": False, "iam": True}, "Invalid connection property setting. SSL must be enabled when using IAM"),
     (
         {"iam": False, "credentials_provider": "anything"},
-        "Invalid connection property setting. IAM must be enabled when using credentials via identity provider",
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": False, "profile": "default"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": False, "access_key_id": "my_key"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": False, "secret_access_key": "shh it's a secret"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": False, "session_token": "my_session"},
+        "Invalid connection property setting",
     ),
     (
         {"iam": True, "ssl": True},
-        "Invalid connection property setting. Credentials provider cannot be None when IAM is enabled",
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "access_key_id": "my_key", "credentials_provider": "OktaCredentialsProvider"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "secret_access_key": "my_secret", "credentials_provider": "OktaCredentialsProvider"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "session_token": "token", "credentials_provider": "OktaCredentialsProvider"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "profile": "default", "credentials_provider": "OktaCredentialsProvider"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "profile": "default", "access_key_id": "my_key"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "profile": "default", "secret_access_key": "my_secret"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "profile": "default", "session_token": "token"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "secret_access_key": "my_secret"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "session_token": "token"},
+        "Invalid connection property setting",
+    ),
+    (
+        {"iam": True, "ssl": True, "access_key_id": "my_key", "password": ""},
+        "Invalid connection property setting",
     ),
     (
         {"iam": False, "ssl_insecure": False},
-        "Invalid connection property setting. IAM must be enabled when using ssl_insecure",
+        "Invalid connection property setting",
     ),
     (
         {"iam": False, "ssl_insecure": True},
-        "Invalid connection property setting. IAM must be enabled when using ssl_insecure",
+        "Invalid connection property setting",
     ),
 ]
 
@@ -213,3 +274,50 @@ def test_set_iam_properties_provider_assigned(mocker, provider):
     assert spy.call_count == 1
     # ensure call to add_Parameter was made on the expected Provider class
     assert isinstance(spy.call_args[0][0], expectedProvider) is True
+
+
+valid_aws_credential_args: typing.List[typing.Dict[str, str]] = [
+    {"profile": "default"},
+    {"access_key_id": "myAccessKey", "secret_access_key": "mySecret"},
+    {"access_key_id": "myAccessKey", "password": "myHiddenSecret"},
+    {"access_key_id": "myAccessKey", "secret_access_key": "mySecret", "session_token": "mySession"},
+]
+
+
+@pytest.mark.parametrize("test_input", valid_aws_credential_args)
+def test_set_iam_properties_via_aws_credentials(mocker, test_input):
+    # spy = mocker.spy("redshift_connector", "set_iam_credentials")
+    info_obj: typing.Dict[str, typing.Any] = get_set_iam_properties_args(**test_input)
+    info_obj["ssl"] = True
+    info_obj["iam"] = True
+
+    mocker.patch("redshift_connector.iam_helper.set_iam_credentials", return_value=None)
+    set_iam_properties(**info_obj)
+
+    for aws_cred_key, aws_cred_val in enumerate(test_input):
+        if aws_cred_key == "profile":
+            assert info_obj["info"].profile == aws_cred_val
+        if aws_cred_key == "access_key_id":
+            assert info_obj["info"].access_key_id == aws_cred_val
+        if aws_cred_key == "secret_access_key":
+            assert info_obj["info"].secret_access_key == aws_cred_val
+        if aws_cred_key == "password":
+            assert info_obj["info"].password == aws_cred_val
+        if aws_cred_key == "session_token":
+            assert info_obj["info"].session_token == aws_cred_val
+
+
+def test_set_iam_credentials_via_aws_credentials(mocker):
+    redshift_property: RedshiftProperty = RedshiftProperty()
+    redshift_property.profile = "profile_val"
+    redshift_property.access_key_id = "access_val"
+    redshift_property.secret_access_key = "secret_val"
+    redshift_property.session_token = "session_val"
+
+    mocker.patch("redshift_connector.iam_helper.set_cluster_credentials", return_value=None)
+    spy = mocker.spy(AWSCredentialsProvider, "add_parameter")
+
+    set_iam_credentials(redshift_property)
+    assert spy.called is True
+    assert spy.call_count == 1
+    assert spy.call_args[0][1] == redshift_property