Skip to content

Commit c5c137b

Browse files
alextwoodsalextwoods
committed
Add basic implementation of Sigv4 AuthScheme + wire up codegen
1 parent f4b87b9 commit c5c137b

File tree

6 files changed

+214
-3
lines changed

6 files changed

+214
-3
lines changed

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

Lines changed: 149 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,159 @@
44
*/
55
package software.amazon.smithy.python.aws.codegen;
66

7+
import java.util.Collections;
8+
import java.util.List;
9+
import software.amazon.smithy.aws.traits.auth.SigV4Trait;
10+
import software.amazon.smithy.codegen.core.Symbol;
11+
import software.amazon.smithy.model.shapes.ShapeId;
12+
import software.amazon.smithy.model.traits.HttpApiKeyAuthTrait;
13+
import software.amazon.smithy.python.codegen.ApplicationProtocol;
14+
import software.amazon.smithy.python.codegen.CodegenUtils;
15+
import software.amazon.smithy.python.codegen.ConfigProperty;
16+
import software.amazon.smithy.python.codegen.DerivedProperty;
17+
import software.amazon.smithy.python.codegen.GenerationContext;
18+
import software.amazon.smithy.python.codegen.SmithyPythonDependency;
19+
import software.amazon.smithy.python.codegen.integrations.AuthScheme;
720
import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
21+
import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin;
822
import software.amazon.smithy.utils.SmithyInternalApi;
923

1024
/**
1125
* Adds support for AWS auth traits.
1226
*/
1327
@SmithyInternalApi
14-
public class AwsAuthIntegration implements PythonIntegration {}
28+
public class AwsAuthIntegration implements PythonIntegration {
29+
private static final String SIGV4_OPTION_GENERATOR_NAME = "_generate_sigv4_option";
30+
31+
@Override
32+
public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
33+
var regionConfig = ConfigProperty.builder()
34+
.name("region")
35+
.type(Symbol.builder().name("str").build())
36+
.documentation(" The AWS region to connect to. The configured region is used to "
37+
+ "determine the service endpoint.")
38+
.build();
39+
40+
return List.of(
41+
RuntimeClientPlugin.builder()
42+
.servicePredicate((model, service) -> service.hasTrait(SigV4Trait.class))
43+
.addConfigProperty(ConfigProperty.builder()
44+
// TODO: Naming of this config RE: backwards compatability/migation considerations
45+
.name("aws_credentials_identity_resolver")
46+
.documentation("Resolves AWS Credentials. Required for operations that use Sigv4 Auth.")
47+
.type(Symbol.builder()
48+
.name("IdentityResolver[AWSCredentialIdentity, IdentityProperties]")
49+
.addReference(Symbol.builder()
50+
.addDependency(SmithyPythonDependency.SMITHY_CORE)
51+
.name("IdentityResolver")
52+
.namespace("smithy_core.aio.interfaces.identity", ".")
53+
.build())
54+
.addReference(Symbol.builder()
55+
.addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
56+
.name("AWSCredentialIdentity")
57+
.namespace("smithy_aws_core.identity", ".")
58+
.build())
59+
.addReference(Symbol.builder()
60+
.addDependency(SmithyPythonDependency.SMITHY_CORE)
61+
.name("IdentityProperties")
62+
.namespace("smithy_core.interfaces.identity", ".")
63+
.build())
64+
.build())
65+
// TODO: Initialize with the provider chain?
66+
.nullable(true)
67+
.build())
68+
.addConfigProperty(regionConfig)
69+
.authScheme(new Sigv4AuthScheme())
70+
.build()
71+
);
72+
}
73+
74+
@Override
75+
public void customize(GenerationContext context) {
76+
if (!hasSigV4Auth(context)) {
77+
return;
78+
}
79+
var trait = context.settings().service(context.model()).expectTrait(SigV4Trait.class);
80+
var params = CodegenUtils.getHttpAuthParamsSymbol(context.settings());
81+
var resolver = CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings());
82+
83+
// Add a function that generates the http auth option for api key auth.
84+
// This needs to be generated because there's modeled parameters that
85+
// must be accounted for.
86+
context.writerDelegator().useFileWriter(resolver.getDefinitionFile(), resolver.getNamespace(), writer -> {
87+
writer.addDependency(SmithyPythonDependency.SMITHY_HTTP);
88+
writer.addImport("smithy_http.aio.interfaces.auth", "HTTPAuthOption");
89+
writer.pushState();
90+
91+
writer.write("""
92+
def $1L(auth_params: $2T) -> HTTPAuthOption | None:
93+
return HTTPAuthOption(
94+
scheme_id=$3S,
95+
identity_properties={},
96+
signer_properties={
97+
"service": $4S,
98+
"region": auth_params.region
99+
}
100+
)
101+
""",
102+
SIGV4_OPTION_GENERATOR_NAME,
103+
params,
104+
SigV4Trait.ID.toString(),
105+
trait.getName());
106+
writer.popState();
107+
});
108+
}
109+
110+
private boolean hasSigV4Auth(GenerationContext context) {
111+
var service = context.settings().service(context.model());
112+
return service.hasTrait(SigV4Trait.class);
113+
}
114+
115+
/**
116+
* The AuthScheme representing api key auth.
117+
*/
118+
private static final class Sigv4AuthScheme implements AuthScheme {
119+
120+
@Override
121+
public ShapeId getAuthTrait() {
122+
return SigV4Trait.ID;
123+
}
124+
125+
@Override
126+
public ApplicationProtocol getApplicationProtocol() {
127+
return ApplicationProtocol.createDefaultHttpApplicationProtocol();
128+
}
129+
130+
@Override
131+
public List<DerivedProperty> getAuthProperties() {
132+
return List.of(
133+
DerivedProperty.builder()
134+
.name("region")
135+
.source(DerivedProperty.Source.CONFIG)
136+
.type(Symbol.builder().name("str").build())
137+
.sourcePropertyName("region")
138+
.build()
139+
);
140+
}
141+
142+
143+
@Override
144+
public Symbol getAuthOptionGenerator(GenerationContext context) {
145+
var resolver = CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings());
146+
return Symbol.builder()
147+
.name(SIGV4_OPTION_GENERATOR_NAME)
148+
.namespace(resolver.getNamespace(), ".")
149+
.definitionFile(resolver.getDefinitionFile())
150+
.build();
151+
}
152+
153+
@Override
154+
public Symbol getAuthSchemeSymbol(GenerationContext context) {
155+
return Symbol.builder()
156+
.name("SigV4AuthScheme")
157+
.namespace("smithy_aws_core.auth.sigv4", ".")
158+
.addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
159+
.build();
160+
}
161+
}
162+
}

‎packages/smithy-aws-core/pyproject.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ readme = "README.md"
66
requires-python = ">=3.12"
77
dependencies = [
88
"smithy-core",
9+
"smithy-http",
10+
"aws-sdk-signers"
911
]
1012

1113
[build-system]

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
from dataclasses import dataclass
4+
from typing import Protocol
5+
6+
from smithy_aws_core.identity import AWSCredentialIdentity
7+
from smithy_core.aio.interfaces.identity import IdentityResolver
8+
from smithy_core.exceptions import SmithyIdentityException
9+
from smithy_core.interfaces.identity import IdentityProperties
10+
from smithy_http.aio.interfaces.auth import HTTPAuthScheme, HTTPSigner
11+
from aws_sdk_signers import SigV4SigningProperties, SigV4Signer
12+
13+
14+
class SigV4Config(Protocol):
15+
aws_credentials_identity_resolver: (
16+
IdentityResolver[AWSCredentialIdentity, IdentityProperties] | None
17+
)
18+
19+
20+
@dataclass(init=False)
21+
class SigV4AuthScheme(
22+
HTTPAuthScheme[
23+
AWSCredentialIdentity, SigV4Config, IdentityProperties, SigV4SigningProperties
24+
]
25+
):
26+
"""SigV4 AuthScheme."""
27+
28+
scheme_id: str
29+
signer: HTTPSigner[AWSCredentialIdentity, SigV4SigningProperties]
30+
31+
def __init__(
32+
self,
33+
*,
34+
signer: HTTPSigner[AWSCredentialIdentity, SigV4SigningProperties] | None = None,
35+
) -> None:
36+
"""Constructor.
37+
38+
:param identity_resolver: The identity resolver to extract the api key identity.
39+
:param signer: The signer used to sign the request.
40+
"""
41+
self.scheme_id = "aws.auth#sigv4"
42+
# TODO: There are type mismatches in the signature of the "sign" method.
43+
self.signer = signer or SigV4Signer() # type: ignore
44+
45+
def identity_resolver(
46+
self, *, config: SigV4Config
47+
) -> IdentityResolver[AWSCredentialIdentity, IdentityProperties]:
48+
if not config.aws_credentials_identity_resolver:
49+
raise SmithyIdentityException(
50+
"Attempted to use SigV4 auth, but aws_credentials_identity_resolver was not "
51+
"set on the config."
52+
)
53+
return config.aws_credentials_identity_resolver

‎packages/smithy-http/src/smithy_http/aio/auth/apikey.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ def identity_resolver(
7474
) -> IdentityResolver[ApiKeyIdentity, IdentityProperties]:
7575
if not config.api_key_identity_resolver:
7676
raise SmithyIdentityException(
77-
"Attempted to use API key auth, but api_key_identity_resolver was not"
77+
"Attempted to use API key auth, but api_key_identity_resolver was not "
7878
"set on the config."
7979
)
8080
return config.api_key_identity_resolver

‎uv.lock‎

Lines changed: 7 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)