Sigv4 AuthScheme + Static/Environment Credentials Providers

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAuthIntegration.java

@@ -4,11 +4,159 @@
  */
 package software.amazon.smithy.python.aws.codegen;
 
+import java.util.Collections;
+import java.util.List;
+import software.amazon.smithy.aws.traits.auth.SigV4Trait;
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.traits.HttpApiKeyAuthTrait;
+import software.amazon.smithy.python.codegen.ApplicationProtocol;
+import software.amazon.smithy.python.codegen.CodegenUtils;
+import software.amazon.smithy.python.codegen.ConfigProperty;
+import software.amazon.smithy.python.codegen.DerivedProperty;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.SmithyPythonDependency;
+import software.amazon.smithy.python.codegen.integrations.AuthScheme;
 import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
+import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
 /**
  * Adds support for AWS auth traits.
  */
 @SmithyInternalApi
-public class AwsAuthIntegration implements PythonIntegration {}
+public class AwsAuthIntegration implements PythonIntegration {
+    private static final String SIGV4_OPTION_GENERATOR_NAME = "_generate_sigv4_option";
+
+    @Override
+    public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
+        var regionConfig = ConfigProperty.builder()
+                .name("region")
+                .type(Symbol.builder().name("str").build())
+                .documentation(" The AWS region to connect to. The configured region is used to "
+                        + "determine the service endpoint.")
+                .build();
+
+        return List.of(
+                RuntimeClientPlugin.builder()
+                        .servicePredicate((model, service) -> service.hasTrait(SigV4Trait.class))
+                        .addConfigProperty(ConfigProperty.builder()
+                                // TODO: Naming of this config RE: backwards compatability/migation considerations
+                                .name("aws_credentials_identity_resolver")
+                                .documentation("Resolves AWS Credentials. Required for operations that use Sigv4 Auth.")
+                                .type(Symbol.builder()
+                                        .name("IdentityResolver[AWSCredentialsIdentity, IdentityProperties]")
+                                        .addReference(Symbol.builder()
+                                                .addDependency(SmithyPythonDependency.SMITHY_CORE)
+                                                .name("IdentityResolver")
+                                                .namespace("smithy_core.aio.interfaces.identity", ".")
+                                                .build())
+                                        .addReference(Symbol.builder()
+                                                .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
+                                                .name("AWSCredentialsIdentity")
+                                                .namespace("smithy_aws_core.identity", ".")
+                                                .build())
+                                        .addReference(Symbol.builder()
+                                                .addDependency(SmithyPythonDependency.SMITHY_CORE)
+                                                .name("IdentityProperties")
+                                                .namespace("smithy_core.interfaces.identity", ".")
+                                                .build())
+                                        .build())
+                                // TODO: Initialize with the provider chain?
+                                .nullable(true)
+                                .build())
+                        .addConfigProperty(regionConfig)
+                        .authScheme(new Sigv4AuthScheme())
+                        .build()
+        );
+    }
+
+    @Override
+    public void customize(GenerationContext context) {
+        if (!hasSigV4Auth(context)) {
+            return;
+        }
+        var trait = context.settings().service(context.model()).expectTrait(SigV4Trait.class);
+        var params = CodegenUtils.getHttpAuthParamsSymbol(context.settings());
+        var resolver = CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings());
+
+        // Add a function that generates the http auth option for api key auth.
+        // This needs to be generated because there's modeled parameters that
+        // must be accounted for.
+        context.writerDelegator().useFileWriter(resolver.getDefinitionFile(), resolver.getNamespace(), writer -> {
+            writer.addDependency(SmithyPythonDependency.SMITHY_HTTP);
+            writer.addImport("smithy_http.aio.interfaces.auth", "HTTPAuthOption");
+            writer.pushState();
+
+            writer.write("""
+                    def $1L(auth_params: $2T) -> HTTPAuthOption | None:
+                        return HTTPAuthOption(
+                            scheme_id=$3S,
+                            identity_properties={},
+                            signer_properties={
+                                "service": $4S,
+                                "region": auth_params.region
+                            }
+                        )
+                    """,
+                    SIGV4_OPTION_GENERATOR_NAME,
+                    params,
+                    SigV4Trait.ID.toString(),
+                    trait.getName());
+            writer.popState();
+        });
+    }
+
+    private boolean hasSigV4Auth(GenerationContext context) {
+        var service = context.settings().service(context.model());
+        return service.hasTrait(SigV4Trait.class);
+    }
+
+    /**
+     * The AuthScheme representing api key auth.
+     */
+    private static final class Sigv4AuthScheme implements AuthScheme {
+
+        @Override
+        public ShapeId getAuthTrait() {
+            return SigV4Trait.ID;
+        }
+
+        @Override
+        public ApplicationProtocol getApplicationProtocol() {
+            return ApplicationProtocol.createDefaultHttpApplicationProtocol();
+        }
+
+        @Override
+        public List<DerivedProperty> getAuthProperties() {
+            return List.of(
+                    DerivedProperty.builder()
+                            .name("region")
+                            .source(DerivedProperty.Source.CONFIG)
+                            .type(Symbol.builder().name("str").build())
+                            .sourcePropertyName("region")
+                            .build()
+            );
+        }
+
+
+        @Override
+        public Symbol getAuthOptionGenerator(GenerationContext context) {
+            var resolver = CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings());
+            return Symbol.builder()
+                    .name(SIGV4_OPTION_GENERATOR_NAME)
+                    .namespace(resolver.getNamespace(), ".")
+                    .definitionFile(resolver.getDefinitionFile())
+                    .build();
+        }
+
+        @Override
+        public Symbol getAuthSchemeSymbol(GenerationContext context) {
+            return Symbol.builder()
+                    .name("SigV4AuthScheme")
+                    .namespace("smithy_aws_core.auth", ".")
+                    .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
+                    .build();
+        }
+    }
+}

codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java

@@ -417,6 +417,7 @@ async def _handle_attempt(
                             for option in auth_options:
                                 if option.scheme_id in config.http_auth_schemes:
                                     auth_option = option
+                                    break
 
                             signer: HTTPSigner[Any, Any] | None = None
                             identity: Identity | None = None

packages/smithy-aws-core/pyproject.toml

@@ -6,6 +6,8 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "smithy-core",
+    "smithy-http",
+    "aws-sdk-signers"
 ]
 
 [build-system]

packages/smithy-aws-core/src/smithy_aws_core/auth/__init__.py

@@ -0,0 +1,6 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+
+from .sigv4 import SigV4AuthScheme
+
+__all__ = ("SigV4AuthScheme",)

packages/smithy-aws-core/src/smithy_aws_core/auth/sigv4.py

@@ -0,0 +1,53 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from typing import Protocol
+
+from smithy_aws_core.identity import AWSCredentialsIdentity
+from smithy_core.aio.interfaces.identity import IdentityResolver
+from smithy_core.exceptions import SmithyIdentityException
+from smithy_core.interfaces.identity import IdentityProperties
+from smithy_http.aio.interfaces.auth import HTTPAuthScheme, HTTPSigner
+from aws_sdk_signers import SigV4SigningProperties, AsyncSigV4Signer
+
+
+class SigV4Config(Protocol):
+    aws_credentials_identity_resolver: (
+        IdentityResolver[AWSCredentialsIdentity, IdentityProperties] | None
+    )
+
+
+@dataclass(init=False)
+class SigV4AuthScheme(
+    HTTPAuthScheme[
+        AWSCredentialsIdentity, SigV4Config, IdentityProperties, SigV4SigningProperties
+    ]
+):
+    """SigV4 AuthScheme."""
+
+    scheme_id: str = "aws.auth#sigv4"
+    signer: HTTPSigner[AWSCredentialsIdentity, SigV4SigningProperties]
+
+    def __init__(
+        self,
+        *,
+        signer: HTTPSigner[AWSCredentialsIdentity, SigV4SigningProperties]
+        | None = None,
+    ) -> None:
+        """Constructor.
+
+        :param identity_resolver: The identity resolver to extract the api key identity.
+        :param signer: The signer used to sign the request.
+        """
+        # TODO: There are type mismatches in the signature of the "sign" method.
+        self.signer = signer or AsyncSigV4Signer()  # type: ignore
+
+    def identity_resolver(
+        self, *, config: SigV4Config
+    ) -> IdentityResolver[AWSCredentialsIdentity, IdentityProperties]:
+        if not config.aws_credentials_identity_resolver:
+            raise SmithyIdentityException(
+                "Attempted to use SigV4 auth, but aws_credentials_identity_resolver was not "
+                "set on the config."
+            )
+        return config.aws_credentials_identity_resolver

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/__init__.py

@@ -0,0 +1,6 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+from .environment import EnvironmentCredentialsResolver
+from .static import StaticCredentialsResolver
+
+__all__ = ("EnvironmentCredentialsResolver", "StaticCredentialsResolver")

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/environment.py

@@ -0,0 +1,42 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+import os
+
+from smithy_aws_core.identity import AWSCredentialsIdentity
+from smithy_core.aio.interfaces.identity import IdentityResolver
+from smithy_core.exceptions import SmithyIdentityException
+from smithy_core.interfaces.identity import IdentityProperties
+
+
+class EnvironmentCredentialsResolver(
+    IdentityResolver[AWSCredentialsIdentity, IdentityProperties]
+):
+    """Resolves AWS Credentials from system environment variables."""
+
+    def __init__(self):
+        self._credentials = None
+
+    async def get_identity(
+        self, *, identity_properties: IdentityProperties
+    ) -> AWSCredentialsIdentity:
+        if self._credentials is not None:
+            return self._credentials
+
+        access_key_id = os.getenv("AWS_ACCESS_KEY_ID")
+        secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY")
+        session_token = os.getenv("AWS_SESSION_TOKEN")
+        account_id = os.getenv("AWS_ACCOUNT_ID")
+
+        if access_key_id is None or secret_access_key is None:
+            raise SmithyIdentityException(
+                "AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY are required"
+            )
+
+        self._credentials = AWSCredentialsIdentity(
+            access_key_id=access_key_id,
+            secret_access_key=secret_access_key,
+            session_token=session_token,
+            account_id=account_id,
+        )
+
+        return self._credentials

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/static.py

@@ -0,0 +1,19 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+from smithy_aws_core.identity import AWSCredentialsIdentity
+from smithy_core.aio.interfaces.identity import IdentityResolver
+from smithy_core.interfaces.identity import IdentityProperties
+
+
+class StaticCredentialsResolver(
+    IdentityResolver[AWSCredentialsIdentity, IdentityProperties]
+):
+    """Resolve Static AWS Credentials."""
+
+    def __init__(self, *, credentials: AWSCredentialsIdentity) -> None:
+        self._credentials = credentials
+
+    async def get_identity(
+        self, *, identity_properties: IdentityProperties
+    ) -> AWSCredentialsIdentity:
+        return self._credentials

packages/smithy-aws-core/src/smithy_aws_core/identity.py

@@ -12,10 +12,12 @@
 # language governing permissions and limitations under the License.
 from datetime import datetime
 
+from smithy_core.aio.interfaces.identity import IdentityResolver
 from smithy_core.identity import Identity
+from smithy_core.interfaces.identity import IdentityProperties
 
 
-class AWSCredentialIdentity(Identity):
+class AWSCredentialsIdentity(Identity):
     """Container for AWS authentication credentials."""
 
     def __init__(
@@ -25,6 +27,7 @@ def __init__(
         secret_access_key: str,
         session_token: str | None = None,
         expiration: datetime | None = None,
+        account_id: str | None = None,
     ) -> None:
         """Initialize the AWSCredentialIdentity.
 
@@ -35,11 +38,13 @@ def __init__(
             the supplied credentials.
         :param expiration: The expiration time of the identity. If time zone is
             provided, it is updated to UTC. The value must always be in UTC.
+        :param account_id: The AWS account's ID.
         """
         super().__init__(expiration=expiration)
         self._access_key_id: str = access_key_id
         self._secret_access_key: str = secret_access_key
         self._session_token: str | None = session_token
+        self._account_id: str | None = account_id
 
     @property
     def access_key_id(self) -> str:
@@ -52,3 +57,12 @@ def secret_access_key(self) -> str:
     @property
     def session_token(self) -> str | None:
         return self._session_token
+
+    @property
+    def account_id(self) -> str | None:
+        return self._account_id
+
+
+type AWSCredentialsResolver = IdentityResolver[
+    AWSCredentialsIdentity, IdentityProperties
+]