diff --git a/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAuthIntegration.java b/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAuthIntegration.java index 07d0ff0f2..44d5fcbed 100644 --- a/codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAuthIntegration.java +++ b/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 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 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(); + } + } +} diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java index 2cdd09717..0cb61438a 100644 --- a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java +++ b/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 diff --git a/packages/smithy-aws-core/pyproject.toml b/packages/smithy-aws-core/pyproject.toml index c23564f28..dc9ef220e 100644 --- a/packages/smithy-aws-core/pyproject.toml +++ b/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] diff --git a/packages/smithy-aws-core/src/smithy_aws_core/auth/__init__.py b/packages/smithy-aws-core/src/smithy_aws_core/auth/__init__.py new file mode 100644 index 000000000..0e19ca81e --- /dev/null +++ b/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",) diff --git a/packages/smithy-aws-core/src/smithy_aws_core/auth/sigv4.py b/packages/smithy-aws-core/src/smithy_aws_core/auth/sigv4.py new file mode 100644 index 000000000..fff9abb79 --- /dev/null +++ b/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 diff --git a/packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/__init__.py b/packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/__init__.py new file mode 100644 index 000000000..3aead11b3 --- /dev/null +++ b/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") diff --git a/packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/environment.py b/packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/environment.py new file mode 100644 index 000000000..a04d93b8d --- /dev/null +++ b/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 diff --git a/packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/static.py b/packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/static.py new file mode 100644 index 000000000..3a2ba9421 --- /dev/null +++ b/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 diff --git a/packages/smithy-aws-core/src/smithy_aws_core/identity.py b/packages/smithy-aws-core/src/smithy_aws_core/identity.py index 67d97b803..638322a7c 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/identity.py +++ b/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 +] diff --git a/packages/smithy-aws-core/tests/unit/credentials_resolvers/test_environment_credentials_resolver.py b/packages/smithy-aws-core/tests/unit/credentials_resolvers/test_environment_credentials_resolver.py new file mode 100644 index 000000000..3d38a8ac5 --- /dev/null +++ b/packages/smithy-aws-core/tests/unit/credentials_resolvers/test_environment_credentials_resolver.py @@ -0,0 +1,68 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from smithy_aws_core.credentials_resolvers import EnvironmentCredentialsResolver +from smithy_core.exceptions import SmithyIdentityException +from smithy_core.interfaces.identity import IdentityProperties + + +async def test_no_values_set(): + with pytest.raises(SmithyIdentityException): + await EnvironmentCredentialsResolver().get_identity( + identity_properties=IdentityProperties() + ) + + +async def test_required_values_missing(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_ACCOUNT_ID", "123456789012") + + with pytest.raises(SmithyIdentityException): + await EnvironmentCredentialsResolver().get_identity( + identity_properties=IdentityProperties() + ) + + +async def test_akid_missing(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret") + + with pytest.raises(SmithyIdentityException): + await EnvironmentCredentialsResolver().get_identity( + identity_properties=IdentityProperties() + ) + + +async def test_secret_missing(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "akid") + + with pytest.raises(SmithyIdentityException): + await EnvironmentCredentialsResolver().get_identity( + identity_properties=IdentityProperties() + ) + + +async def test_minimum_required(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "akid") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret") + + credentials = await EnvironmentCredentialsResolver().get_identity( + identity_properties=IdentityProperties() + ) + assert credentials.access_key_id == "akid" + assert credentials.secret_access_key == "secret" + + +async def test_all_values(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "akid") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret") + monkeypatch.setenv("AWS_SESSION_TOKEN", "session") + monkeypatch.setenv("AWS_ACCOUNT_ID", "123456789012") + + credentials = await EnvironmentCredentialsResolver().get_identity( + identity_properties=IdentityProperties() + ) + assert credentials.access_key_id == "akid" + assert credentials.secret_access_key == "secret" + assert credentials.session_token == "session" + assert credentials.account_id == "123456789012" diff --git a/packages/smithy-http/src/smithy_http/aio/auth/apikey.py b/packages/smithy-http/src/smithy_http/aio/auth/apikey.py index cacb81b0c..592f872cc 100644 --- a/packages/smithy-http/src/smithy_http/aio/auth/apikey.py +++ b/packages/smithy-http/src/smithy_http/aio/auth/apikey.py @@ -74,7 +74,7 @@ def identity_resolver( ) -> IdentityResolver[ApiKeyIdentity, IdentityProperties]: if not config.api_key_identity_resolver: raise SmithyIdentityException( - "Attempted to use API key auth, but api_key_identity_resolver was not" + "Attempted to use API key auth, but api_key_identity_resolver was not " "set on the config." ) return config.api_key_identity_resolver diff --git a/uv.lock b/uv.lock index 9768bb114..729168555 100644 --- a/uv.lock +++ b/uv.lock @@ -653,11 +653,17 @@ name = "smithy-aws-core" version = "0.0.1" source = { editable = "packages/smithy-aws-core" } dependencies = [ + { name = "aws-sdk-signers" }, { name = "smithy-core" }, + { name = "smithy-http" }, ] [package.metadata] -requires-dist = [{ name = "smithy-core", editable = "packages/smithy-core" }] +requires-dist = [ + { name = "aws-sdk-signers", editable = "packages/aws-sdk-signers" }, + { name = "smithy-core", editable = "packages/smithy-core" }, + { name = "smithy-http", editable = "packages/smithy-http" }, +] [[package]] name = "smithy-core"