Skip to content

Commit 6023a9a

Browse files
authored
STS: integrate with core response serializer (#9107)
1 parent d37a436 commit 6023a9a

File tree

3 files changed

+77
-152
lines changed

3 files changed

+77
-152
lines changed

moto/sts/exceptions.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
from typing import Any
22

3-
from moto.core.exceptions import RESTError
3+
from moto.core.exceptions import ServiceException
44

55

6-
class STSClientError(RESTError):
7-
code = 400
6+
class STSClientError(ServiceException):
7+
pass
88

99

1010
class STSValidationError(STSClientError):

moto/sts/models.py

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
from moto.core.base_backend import BackendDict, BaseBackend
99
from moto.core.common_models import BaseModel
10-
from moto.core.utils import iso_8601_datetime_with_milliseconds, utcnow
10+
from moto.core.utils import utcnow
1111
from moto.iam.models import AccessKey, iam_backends
1212
from moto.sts.utils import (
1313
DEFAULT_STS_SESSION_DURATION,
@@ -24,10 +24,6 @@ def __init__(self, duration: int, name: Optional[str] = None):
2424
self.name = name
2525
self.policy = None
2626

27-
@property
28-
def expiration_ISO8601(self) -> str:
29-
return iso_8601_datetime_with_milliseconds(self.expiration)
30-
3127

3228
class AssumedRole(BaseModel):
3329
def __init__(
@@ -55,10 +51,6 @@ def __init__(
5551
self.session_token = random_session_token()
5652
self.partition = get_partition(region_name)
5753

58-
@property
59-
def expiration_ISO8601(self) -> str:
60-
return iso_8601_datetime_with_milliseconds(self.expiration)
61-
6254
@property
6355
def user_id(self) -> str:
6456
iam_backend = iam_backends[self.account_id][self.partition]

moto/sts/responses.py

Lines changed: 73 additions & 140 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from moto.core.responses import BaseResponse
1+
from moto.core.responses import ActionResult, BaseResponse
22

33
from .exceptions import STSValidationError
44
from .models import STSBackend, sts_backends
@@ -19,13 +19,20 @@ def _determine_resource(self) -> str:
1919
return self.querystring.get("RoleArn")[0] # type: ignore[index]
2020
return "*"
2121

22-
def get_session_token(self) -> str:
22+
def get_session_token(self) -> ActionResult:
2323
duration = int(self.querystring.get("DurationSeconds", [43200])[0])
2424
token = self.backend.get_session_token(duration=duration)
25-
template = self.response_template(GET_SESSION_TOKEN_RESPONSE)
26-
return template.render(token=token)
27-
28-
def get_federation_token(self) -> str:
25+
result = {
26+
"Credentials": {
27+
"AccessKeyId": "AKIAIOSFODNN7EXAMPLE",
28+
"SecretAccessKey": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLEKEY",
29+
"SessionToken": "AQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKwRcOIfrRh3c/LTo6UDdyJwOOvEVPvLXCrrrUtdnniCEXAMPLE/IvU1dYUg2RVAJBanLiHb4IgRmpRV3zrkuWJOgQs8IZZaIv2BXIa2R4OlgkBN9bkUDNCJiBeb/AXlzBBko7b15fjrBs2+cTQtpZ3CYWFXG8C5zqx37wnOE49mRl/+OtkIKGO7fAE",
30+
"Expiration": token.expiration,
31+
}
32+
}
33+
return ActionResult(result)
34+
35+
def get_federation_token(self) -> ActionResult:
2936
duration = int(self.querystring.get("DurationSeconds", [43200])[0])
3037
policy = self.querystring.get("Policy", [None])[0]
3138

@@ -39,12 +46,22 @@ def get_federation_token(self) -> str:
3946

4047
name = self.querystring.get("Name")[0] # type: ignore
4148
token = self.backend.get_federation_token(duration=duration, name=name)
42-
template = self.response_template(GET_FEDERATION_TOKEN_RESPONSE)
43-
return template.render(
44-
token=token, account_id=self.current_account, partition=self.partition
45-
)
46-
47-
def assume_role(self) -> str:
49+
result = {
50+
"Credentials": {
51+
"AccessKeyId": "AKIAIOSFODNN7EXAMPLE",
52+
"SecretAccessKey": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLEKEY",
53+
"SessionToken": "AQoDYXdzEPT//////////wEXAMPLEtc764bNrC9SAPBSM22wDOk4x4HIZ8j4FZTwdQWLWsKWHGBuFqwAeMicRXmxfpSPfIeoIYRqTflfKD8YUuwthAx7mSEI/qkPpKPi/kMcGdQrmGdeehM4IC1NtBmUpp2wUE8phUZampKsburEDy0KPkyQDYwT7WZ0wq5VSXDvp75YU9HFvlRd8Tx6q6fE8YQcHNVXAkiY9q6d+xo0rKwT38xVqr7ZD0u0iPPkUL64lIZbqBAz+scqKmlzm8FDrypNC9Yjc8fPOLn9FX9KSYvKTr4rvx3iSIlTJabIQwj2ICCR/oLxBA==",
54+
"Expiration": token.expiration,
55+
},
56+
"FederatedUser": {
57+
"FederatedUserId": f"{self.current_account}:{token.name}",
58+
"Arn": f"arn:{self.partition}:sts::{self.current_account}:federated-user/{token.name}",
59+
},
60+
"PackedPolicySize": 6,
61+
}
62+
return ActionResult(result)
63+
64+
def assume_role(self) -> ActionResult:
4865
role_session_name = self.querystring.get("RoleSessionName")[0] # type: ignore
4966
role_arn = self.querystring.get("RoleArn")[0] # type: ignore
5067

@@ -60,10 +77,17 @@ def assume_role(self) -> str:
6077
duration=duration,
6178
external_id=external_id,
6279
)
63-
template = self.response_template(ASSUME_ROLE_RESPONSE)
64-
return template.render(role=role)
65-
66-
def assume_role_with_web_identity(self) -> str:
80+
result = {
81+
"Credentials": role,
82+
"AssumedRoleUser": {
83+
"AssumedRoleId": role.user_id,
84+
"Arn": role.arn,
85+
},
86+
"PackedPolicySize": 6,
87+
}
88+
return ActionResult(result)
89+
90+
def assume_role_with_web_identity(self) -> ActionResult:
6791
role_session_name = self.querystring.get("RoleSessionName")[0] # type: ignore
6892
role_arn = self.querystring.get("RoleArn")[0] # type: ignore
6993

@@ -79,10 +103,17 @@ def assume_role_with_web_identity(self) -> str:
79103
duration=duration,
80104
external_id=external_id,
81105
)
82-
template = self.response_template(ASSUME_ROLE_WITH_WEB_IDENTITY_RESPONSE)
83-
return template.render(role=role)
84-
85-
def assume_role_with_saml(self) -> str:
106+
result = {
107+
"Credentials": role,
108+
"AssumedRoleUser": {
109+
"AssumedRoleId": f"ARO123EXAMPLE123:{role.session_name}",
110+
"Arn": role.arn,
111+
},
112+
"PackedPolicySize": 6,
113+
}
114+
return ActionResult(result)
115+
116+
def assume_role_with_saml(self) -> ActionResult:
86117
role_arn = self.querystring.get("RoleArn")[0] # type: ignore
87118
principal_arn = self.querystring.get("PrincipalArn")[0] # type: ignore
88119
saml_assertion = self.querystring.get("SAMLAssertion")[0] # type: ignore
@@ -92,127 +123,29 @@ def assume_role_with_saml(self) -> str:
92123
principal_arn=principal_arn,
93124
saml_assertion=saml_assertion,
94125
)
95-
template = self.response_template(ASSUME_ROLE_WITH_SAML_RESPONSE)
96-
return template.render(role=role)
97-
98-
def get_caller_identity(self) -> str:
99-
template = self.response_template(GET_CALLER_IDENTITY_RESPONSE)
100-
126+
result = {
127+
"Credentials": role,
128+
"AssumedRoleUser": {
129+
"AssumedRoleId": role.user_id,
130+
"Arn": role.arn,
131+
},
132+
"PackedPolicySize": 123,
133+
"Subject": role.user_id,
134+
"SubjectType": "persistent",
135+
"Issuer": "http://localhost:3000/",
136+
"Audience": "https://signin.aws.amazon.com/saml",
137+
"NameQualifier": "B64EncodedStringOfHashOfIssuerAccountIdAndUserId=",
138+
}
139+
return ActionResult(result)
140+
141+
def get_caller_identity(self) -> ActionResult:
101142
access_key_id = self.get_access_key()
102143
user_id, arn, account_id = self.backend.get_caller_identity(
103144
access_key_id, self.region
104145
)
105-
106-
return template.render(account_id=account_id, user_id=user_id, arn=arn)
107-
108-
109-
GET_SESSION_TOKEN_RESPONSE = """<GetSessionTokenResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
110-
<GetSessionTokenResult>
111-
<Credentials>
112-
<SessionToken>AQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKwRcOIfrRh3c/LTo6UDdyJwOOvEVPvLXCrrrUtdnniCEXAMPLE/IvU1dYUg2RVAJBanLiHb4IgRmpRV3zrkuWJOgQs8IZZaIv2BXIa2R4OlgkBN9bkUDNCJiBeb/AXlzBBko7b15fjrBs2+cTQtpZ3CYWFXG8C5zqx37wnOE49mRl/+OtkIKGO7fAE</SessionToken>
113-
<SecretAccessKey>wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLEKEY</SecretAccessKey>
114-
<Expiration>{{ token.expiration_ISO8601 }}</Expiration>
115-
<AccessKeyId>AKIAIOSFODNN7EXAMPLE</AccessKeyId>
116-
</Credentials>
117-
</GetSessionTokenResult>
118-
<ResponseMetadata>
119-
<RequestId>58c5dbae-abef-11e0-8cfe-09039844ac7d</RequestId>
120-
</ResponseMetadata>
121-
</GetSessionTokenResponse>"""
122-
123-
124-
GET_FEDERATION_TOKEN_RESPONSE = """<GetFederationTokenResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
125-
<GetFederationTokenResult>
126-
<Credentials>
127-
<SessionToken>AQoDYXdzEPT//////////wEXAMPLEtc764bNrC9SAPBSM22wDOk4x4HIZ8j4FZTwdQWLWsKWHGBuFqwAeMicRXmxfpSPfIeoIYRqTflfKD8YUuwthAx7mSEI/qkPpKPi/kMcGdQrmGdeehM4IC1NtBmUpp2wUE8phUZampKsburEDy0KPkyQDYwT7WZ0wq5VSXDvp75YU9HFvlRd8Tx6q6fE8YQcHNVXAkiY9q6d+xo0rKwT38xVqr7ZD0u0iPPkUL64lIZbqBAz+scqKmlzm8FDrypNC9Yjc8fPOLn9FX9KSYvKTr4rvx3iSIlTJabIQwj2ICCR/oLxBA==</SessionToken>
128-
<SecretAccessKey>wJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLEKEY</SecretAccessKey>
129-
<Expiration>{{ token.expiration_ISO8601 }}</Expiration>
130-
<AccessKeyId>AKIAIOSFODNN7EXAMPLE</AccessKeyId>
131-
</Credentials>
132-
<FederatedUser>
133-
<Arn>arn:{{ partition }}:sts::{{ account_id }}:federated-user/{{ token.name }}</Arn>
134-
<FederatedUserId>{{ account_id }}:{{ token.name }}</FederatedUserId>
135-
</FederatedUser>
136-
<PackedPolicySize>6</PackedPolicySize>
137-
</GetFederationTokenResult>
138-
<ResponseMetadata>
139-
<RequestId>c6104cbe-af31-11e0-8154-cbc7ccf896c7</RequestId>
140-
</ResponseMetadata>
141-
</GetFederationTokenResponse>"""
142-
143-
144-
ASSUME_ROLE_RESPONSE = """<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
145-
<AssumeRoleResult>
146-
<Credentials>
147-
<SessionToken>{{ role.session_token }}</SessionToken>
148-
<SecretAccessKey>{{ role.secret_access_key }}</SecretAccessKey>
149-
<Expiration>{{ role.expiration_ISO8601 }}</Expiration>
150-
<AccessKeyId>{{ role.access_key_id }}</AccessKeyId>
151-
</Credentials>
152-
<AssumedRoleUser>
153-
<Arn>{{ role.arn }}</Arn>
154-
<AssumedRoleId>{{ role.user_id }}</AssumedRoleId>
155-
</AssumedRoleUser>
156-
<PackedPolicySize>6</PackedPolicySize>
157-
</AssumeRoleResult>
158-
<ResponseMetadata>
159-
<RequestId>c6104cbe-af31-11e0-8154-cbc7ccf896c7</RequestId>
160-
</ResponseMetadata>
161-
</AssumeRoleResponse>"""
162-
163-
164-
ASSUME_ROLE_WITH_WEB_IDENTITY_RESPONSE = """<AssumeRoleWithWebIdentityResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
165-
<AssumeRoleWithWebIdentityResult>
166-
<Credentials>
167-
<SessionToken>{{ role.session_token }}</SessionToken>
168-
<SecretAccessKey>{{ role.secret_access_key }}</SecretAccessKey>
169-
<Expiration>{{ role.expiration_ISO8601 }}</Expiration>
170-
<AccessKeyId>{{ role.access_key_id }}</AccessKeyId>
171-
</Credentials>
172-
<AssumedRoleUser>
173-
<Arn>{{ role.arn }}</Arn>
174-
<AssumedRoleId>ARO123EXAMPLE123:{{ role.session_name }}</AssumedRoleId>
175-
</AssumedRoleUser>
176-
<PackedPolicySize>6</PackedPolicySize>
177-
</AssumeRoleWithWebIdentityResult>
178-
<ResponseMetadata>
179-
<RequestId>c6104cbe-af31-11e0-8154-cbc7ccf896c7</RequestId>
180-
</ResponseMetadata>
181-
</AssumeRoleWithWebIdentityResponse>"""
182-
183-
184-
ASSUME_ROLE_WITH_SAML_RESPONSE = """<AssumeRoleWithSAMLResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
185-
<AssumeRoleWithSAMLResult>
186-
<Audience>https://signin.aws.amazon.com/saml</Audience>
187-
<AssumedRoleUser>
188-
<AssumedRoleId>{{ role.user_id }}</AssumedRoleId>
189-
<Arn>{{ role.arn }}</Arn>
190-
</AssumedRoleUser>
191-
<Credentials>
192-
<AccessKeyId>{{ role.access_key_id }}</AccessKeyId>
193-
<SecretAccessKey>{{ role.secret_access_key }}</SecretAccessKey>
194-
<SessionToken>{{ role.session_token }}</SessionToken>
195-
<Expiration>{{ role.expiration_ISO8601 }}</Expiration>
196-
</Credentials>
197-
<Subject>{{ role.user_id }}</Subject>
198-
<NameQualifier>B64EncodedStringOfHashOfIssuerAccountIdAndUserId=</NameQualifier>
199-
<SubjectType>persistent</SubjectType>
200-
<Issuer>http://localhost:3000/</Issuer>
201-
</AssumeRoleWithSAMLResult>
202-
<ResponseMetadata>
203-
<RequestId>c6104cbe-af31-11e0-8154-cbc7ccf896c7</RequestId>
204-
</ResponseMetadata>
205-
</AssumeRoleWithSAMLResponse>"""
206-
207-
208-
GET_CALLER_IDENTITY_RESPONSE = """<GetCallerIdentityResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
209-
<GetCallerIdentityResult>
210-
<Arn>{{ arn }}</Arn>
211-
<UserId>{{ user_id }}</UserId>
212-
<Account>{{ account_id }}</Account>
213-
</GetCallerIdentityResult>
214-
<ResponseMetadata>
215-
<RequestId>c6104cbe-af31-11e0-8154-cbc7ccf896c7</RequestId>
216-
</ResponseMetadata>
217-
</GetCallerIdentityResponse>
218-
"""
146+
result = {
147+
"UserId": user_id,
148+
"Account": account_id,
149+
"Arn": arn,
150+
}
151+
return ActionResult(result)

0 commit comments

Comments
 (0)