feat: federated auth plugin (#373)

Author: joyc-bq

aws_advanced_python_wrapper/federated_plugin.py

@@ -16,6 +16,8 @@
 
 from typing import TYPE_CHECKING
 
+from aws_advanced_python_wrapper.utils.iamutils import IamAuthUtils, TokenInfo
+
 if TYPE_CHECKING:
     from boto3 import Session
     from aws_advanced_python_wrapper.driver_dialect import DriverDialect
@@ -39,23 +41,6 @@
 logger = Logger(__name__)
 
 
-class TokenInfo:
-    @property
-    def token(self):
-        return self._token
-
-    @property
-    def expiration(self):
-        return self._expiration
-
-    def __init__(self, token: str, expiration: datetime):
-        self._token = token
-        self._expiration = expiration
-
-    def is_expired(self) -> bool:
-        return datetime.now() > self._expiration
-
-
 class IamAuthPlugin(Plugin):
     _SUBSCRIBED_METHODS: Set[str] = {"connect", "force_connect"}
     # Leave 30 second buffer to prevent time-of-check to time-of-use errors
@@ -86,10 +71,10 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
         if not WrapperProperties.USER.get(props):
             raise AwsWrapperError(Messages.get_formatted("IamPlugin.IsNoneOrEmpty", WrapperProperties.USER.name))
 
-        host = WrapperProperties.IAM_HOST.get(props) if WrapperProperties.IAM_HOST.get(props) else host_info.host
+        host = IamAuthUtils.get_iam_host(props, host_info)
         region = WrapperProperties.IAM_REGION.get(props) \
             if WrapperProperties.IAM_REGION.get(props) else self._get_rds_region(host)
-        port = self._get_port(props, host_info)
+        port = IamAuthUtils.get_port(props, host_info, self._plugin_service.database_dialect.default_port)
         token_expiration_sec: int = WrapperProperties.IAM_EXPIRATION.get_int(props)
 
         cache_key: str = self._get_cache_key(
@@ -170,27 +155,11 @@ def _generate_authentication_token(self,
     def _get_cache_key(self, user: Optional[str], hostname: Optional[str], port: int, region: Optional[str]) -> str:
         return f"{region}:{hostname}:{port}:{user}"
 
-    def _get_port(self, props: Properties, host_info: HostInfo) -> int:
-        if WrapperProperties.IAM_DEFAULT_PORT.get(props):
-            default_port: int = WrapperProperties.IAM_DEFAULT_PORT.get_int(props)
-            if default_port > 0:
-                return default_port
-            else:
-                logger.debug("IamAuthPlugin.InvalidPort", default_port)
-
-        if host_info.is_port_specified():
-            return host_info.port
-
-        if self._plugin_service.database_dialect is not None:
-            return self._plugin_service.database_dialect.default_port
-
-        raise AwsWrapperError(Messages.get("IamAuthPlugin.NoValidPorts"))
-
     def _get_rds_region(self, hostname: Optional[str]) -> str:
         rds_region = self._rds_utils.get_rds_region(hostname) if hostname else None
 
         if not rds_region:
-            exception_message = "IamAuthPlugin.UnsupportedHostname"
+            exception_message = "RdsUtils.UnsupportedHostname"
             logger.debug(exception_message, hostname)
             raise AwsWrapperError(Messages.get_formatted(exception_message, hostname))
 

aws_advanced_python_wrapper/iam_plugin.py

@@ -16,6 +16,9 @@
 
 from typing import TYPE_CHECKING, ClassVar, List, Type
 
+from aws_advanced_python_wrapper.federated_plugin import \
+    FederatedAuthPluginFactory
+
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.driver_dialect import DriverDialect
     from aws_advanced_python_wrapper.driver_dialect_manager import DriverDialectManager
@@ -563,7 +566,8 @@ class PluginManager(CanReleaseResources):
         "stale_dns": StaleDnsPluginFactory,
         "connect_time": ConnectTimePluginFactory,
         "execute_time": ExecuteTimePluginFactory,
-        "dev": DeveloperPluginFactory
+        "dev": DeveloperPluginFactory,
+        "federated_auth": FederatedAuthPluginFactory
     }
 
     WEIGHT_RELATIVE_TO_PRIOR_PLUGIN = -1
@@ -581,7 +585,8 @@ class PluginManager(CanReleaseResources):
         AwsSecretsManagerPluginFactory: 700,
         ConnectTimePluginFactory: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN,
         ExecuteTimePluginFactory: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN,
-        DeveloperPluginFactory: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN
+        DeveloperPluginFactory: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN,
+        FederatedAuthPluginFactory: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN
     }
 
     def __init__(self, container: PluginServiceManagerContainer, props: Properties):

aws_advanced_python_wrapper/plugin_service.py

@@ -17,6 +17,14 @@
 AuroraPgDialect.HasExtensionsTrue=[AuroraPgDialect] has_extensions: True
 AuroraPgDialect.HasTopologyTrue=[AuroraPgDialect] has_topology: True
 
+AdfsCredentialsProviderFactory.FailedLogin=[AdfsCredentialsProviderFactory] Failed login. Could not obtain SAML Assertion from ADFS SignOn Page POST response: '{}'
+AdfsCredentialsProviderFactory.GetSamlAssertionFailed=[AdfsCredentialsProviderFactory] Failed to get SAML Assertion due to exception: '{}'
+AdfsCredentialsProviderFactory.InvalidHttpsUrl=[AdfsCredentialsProviderFactory] Invalid HTTPS URL: '{}'
+AdfsCredentialsProviderFactory.SignOnPagePostActionUrl=[AdfsCredentialsProviderFactory] ADFS SignOn Action URL: '{}'
+AdfsCredentialsProviderFactory.SignOnPagePostActionRequestFailed=[AdfsCredentialsProviderFactory] ADFS SignOn Page POST action failed with HTTP status '{}', reason phrase '{}', and response '{}'
+AdfsCredentialsProviderFactory.SignOnPageRequestFailed=[AdfsCredentialsProviderFactory] ADFS SignOn Page Request Failed with HTTP status '{}', reason phrase '{}', and response '{}'
+AdfsCredentialsProviderFactory.SignOnPageUrl=[AdfsCredentialsProviderFactory] ADFS SignOn URL: '{}'
+
 AwsSdk.UnsupportedRegion=[AwsSdk] Unsupported AWS region {}. For supported regions please read https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/Concepts.RegionsAndAvailabilityZones.html
 
 AwsSecretsManagerPlugin.ConnectException=[AwsSecretsManagerPlugin] Error occurred while opening a connection: {}
@@ -82,6 +90,10 @@ FailoverPlugin.TransactionResolutionUnknownError=[Failover] Transaction resoluti
 FailoverPlugin.UnableToConnectToReader=[Failover] Unable to establish SQL connection to the reader instance.
 FailoverPlugin.UnableToConnectToWriter=[Failover] Unable to establish SQL connection to the writer instance.
 
+FederatedAuthPlugin.UnhandledException=[FederatedAuthPlugin] Unhandled exception: '{}'
+
+FederatedAuthPluginFactory.UnsupportedIdp=[FederatedAuthPluginFactory] Unsupported Identity Provider '{}'. Please visit to the documentation for supported Identity Providers.
+
 HostAvailabilityStrategy.InvalidInitialBackoffTime=[HostAvailabilityStrategy] Invalid value of {} for configuration parameter `host_availability_strategy_initial_backoff_time`. It must be an integer greater than 1.
 HostAvailabilityStrategy.InvalidMaxRetries=[HostAvailabilityStrategy] Invalid value of {} for configuration parameter `host_availability_strategy_max_retries`. It must be an integer greater than 1.
 
@@ -104,7 +116,6 @@ IamAuthPlugin.GeneratedNewIamToken=[IamAuthPlugin] Generated new IAM token = {}
 IamAuthPlugin.InvalidPort=[IamAuthPlugin] Port number: {} is not valid. Port number should be greater than zero. Falling back to default port.
 IamAuthPlugin.NoValidPort=[IamAuthPlugin] Unable to determine a valid port.
 IamAuthPlugin.UnhandledException=[IamAuthPlugin] Unhandled exception: {}
-IamAuthPlugin.UnsupportedHostname=[IamAuthPlugin] Unsupported AWS hostname {}. Amazon domain name in format *.AWS-Region.rds.amazonaws.com or *.rds.AWS-Region.amazonaws.com.cn is expected.
 IamAuthPlugin.UseCachedIamToken=[IamAuthPlugin] Used cached IAM token = {}
 
 IamPlugin.IsNoneOrEmpty=[IamPlugin] Property "{}" is None or empty.
@@ -185,6 +196,8 @@ RdsTestUtility.InvalidDatabaseEngine=[RdsTestUtility] The detected database engi
 RdsTestUtility.MethodNotSupportedForDeployment=[RdsTestUtility] Method '{}' is not supported for the current database engine deployment: '{}'
 RdsTestUtility.WriterInstanceNotFound=[RdsTestUtility] Cannot find writer instance for cluster '{}'.
 
+RdsUtils.UnsupportedHostname=[RdsUtils] Unsupported AWS hostname {}. Amazon domain name in format *.AWS-Region.rds.amazonaws.com or *.rds.AWS-Region.amazonaws.com.cn is expected.
+
 ReaderFailoverHandler.AttemptingReaderConnection=[ReaderFailoverHandler] Trying to connect to reader: '{}', with properties '{}'
 ReaderFailoverHandler.FailedReaderConnection=[ReaderFailoverHandler] Failed to connect to reader: '{}'
 ReaderFailoverHandler.InvalidTopology=[ReaderFailoverHandler] '{}' was called with an invalid (None or empty) topology.

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -0,0 +1,59 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.hostinfo import HostInfo
+
+from aws_advanced_python_wrapper.utils.properties import (Properties,
+                                                          WrapperProperties)
+
+
+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
+
+    @staticmethod
+    def get_port(props: Properties, host_info: HostInfo, dialect_default_port: int) -> int:
+        default_port: int = WrapperProperties.IAM_DEFAULT_PORT.get_int(props)
+        if default_port > 0:
+            return default_port
+
+        if host_info.is_port_specified():
+            return host_info.port
+
+        return dialect_default_port
+
+
+class TokenInfo:
+    @property
+    def token(self):
+        return self._token
+
+    @property
+    def expiration(self):
+        return self._expiration
+
+    def __init__(self, token: str, expiration: datetime):
+        self._token = token
+        self._expiration = expiration
+
+    def is_expired(self) -> bool:
+        return datetime.now() > self._expiration

aws_advanced_python_wrapper/utils/iamutils.py

@@ -234,6 +234,54 @@ class WrapperProperties:
     ROUND_ROBIN_HOST_WEIGHT_PAIRS = WrapperProperty("round_robin_host_weight_pairs",
                                                     "Comma separated list of database host-weight pairs in the format of `<host>:<weight>`.",
                                                     "")
+    # Federated Auth Plugin
+    IDP_ENDPOINT = WrapperProperty("idp_endpoint",
+                                   "The hosting URL of the Identity Provider",
+                                   None)
+
+    IDP_PORT = WrapperProperty("idp_port",
+                               "The hosting port of the Identity Provider",
+                               443)
+
+    RELAYING_PARTY_ID = WrapperProperty("rp_identifier",
+                                        "The relaying party identifier",
+                                        "urn:amazon:webservices")
+
+    IAM_ROLE_ARN = WrapperProperty("iam_role_arn",
+                                   "The ARN of the IAM Role that is to be assumed.",
+                                   None)
+
+    IAM_IDP_ARN = WrapperProperty("iam_idp_arn",
+                                  "The ARN of the Identity Provider",
+                                  None)
+
+    IAM_TOKEN_EXPIRATION = WrapperProperty("iam_token_expiration",
+                                           "IAM token cache expiration in seconds",
+                                           15 * 60 - 30)
+
+    IDP_USERNAME = WrapperProperty("idp_username",
+                                   "The federated user name",
+                                   None)
+
+    IDP_PASSWORD = WrapperProperty("idp_password",
+                                   "The federated user password",
+                                   None)
+
+    HTTP_REQUEST_TIMEOUT = WrapperProperty("http_request_connect_timeout",
+                                           "The timeout value in seconds to send the HTTP request data used by the FederatedAuthPlugin",
+                                           60)
+
+    SSL_SECURE = WrapperProperty("ssl_secure",
+                                 "Whether the SSL session is to be secure and the server's certificates will be verified",
+                                 False)
+
+    IDP_NAME = WrapperProperty("idp_name",
+                               "The name of the Identity Provider implementation used",
+                               "adfs")
+
+    DB_USER = WrapperProperty("db_user",
+                              "The database user used to access the database",
+                              None)
 
 
 class PropertiesUtils:

aws_advanced_python_wrapper/utils/properties.py

@@ -117,7 +117,7 @@ unset FILTER  # Done testing the IAM tests, unset FILTER
 
 | Environment Variable Name | Required | Description                                                                                                                                                                                                      | Example Value                                |
 |---------------------------|----------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------|
-| `DB_USERNAME`             | Yes      | The username to access the database.                                                                                                                                                                             | `admin`                                      |
+| `DB_USER`             | Yes      | The username to access the database.                                                                                                                                                                             | `admin`                                      |
 | `DB_PASSWORD`             | Yes      | The database cluster password.                                                                                                                                                                                   | `password`                                   |
 | `DB_DATABASE_NAME`        | No       | Name of the database that will be used by the tests. The default database name is test.                                                                                                                          | `test_db_name`                               |
 | `AURORA_CLUSTER_NAME`     | Yes      | The database identifier for your Aurora cluster. Must be a unique value to avoid conflicting with existing clusters.                                                                                             | `db-identifier`                              |

docs/development-guide/IntegrationTests.md

@@ -0,0 +1,38 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import mysql.connector
+
+from aws_advanced_python_wrapper import AwsWrapperConnection
+
+if __name__ == "__main__":
+    with AwsWrapperConnection.connect(
+            mysql.connector.Connect,
+            host="database.cluster-xyz.us-east-2.rds.amazonaws.com",
+            database="mysql",
+            plugins="federated_auth",
+            idp_name="adfs",
+            idp_endpoint="ec2amaz-ab3cdef.example.com",
+            iam_role_arn="arn:aws:iam::123456789012:role/adfs_example_iam_role",
+            iam_idp_arn="arn:aws:iam::123456789012:saml-provider/adfs_example",
+            iam_region="us-east-2",
+            idp_username="some_federated_username@example.com",
+            idp_password="some_password",
+            user="john",
+            autocommit=True
+    ) as awsconn, awsconn.cursor() as awscursor:
+        awscursor.execute("SELECT 1")
+
+        res = awscursor.fetchone()
+        print(res)

docs/examples/MySQLFederatedAuthentication.py

@@ -0,0 +1,38 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import psycopg
+
+from aws_advanced_python_wrapper import AwsWrapperConnection
+
+if __name__ == "__main__":
+    with AwsWrapperConnection.connect(
+            psycopg.Connection.connect,
+            host="database.cluster-xyz.us-east-2.rds.amazonaws.com",
+            dbname="postgres",
+            plugins="federated_auth",
+            idp_name="adfs",
+            idp_endpoint="ec2amaz-ab3cdef.example.com",
+            iam_role_arn="arn:aws:iam::123456789012:role/adfs_example_iam_role",
+            iam_idp_arn="arn:aws:iam::123456789012:saml-provider/adfs_example",
+            iam_region="us-east-2",
+            idp_username="some_federated_username@example.com",
+            idp_password="some_password",
+            user="john",
+            autocommit=True
+    ) as awsconn, awsconn.cursor() as awscursor:
+        awscursor.execute("SELECT 1")
+
+        res = awscursor.fetchone()
+        print(res)