186 lines
		
	
	
		
			8.7 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
			
		
		
	
	
			186 lines
		
	
	
		
			8.7 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
import boto3
 | 
						|
import unittest
 | 
						|
 | 
						|
from base64 import b64encode
 | 
						|
from moto import mock_dynamodb, mock_sts, mock_iam
 | 
						|
from moto.core import DEFAULT_ACCOUNT_ID as ACCOUNT_ID
 | 
						|
 | 
						|
 | 
						|
@mock_sts
 | 
						|
@mock_iam
 | 
						|
@mock_dynamodb
 | 
						|
class TestStsAssumeRole(unittest.TestCase):
 | 
						|
    def setUp(self) -> None:
 | 
						|
        self.account_b = "111111111111"
 | 
						|
        self.sts = boto3.client("sts", region_name="us-east-1")
 | 
						|
 | 
						|
    def test_assume_role_in_different_account(self):
 | 
						|
        # assume role to another aws account
 | 
						|
        role_name = f"arn:aws:iam::{self.account_b}:role/my-role"
 | 
						|
        response = self.sts.assume_role(
 | 
						|
            RoleArn=role_name,
 | 
						|
            RoleSessionName="test-session-name",
 | 
						|
            ExternalId="test-external-id",
 | 
						|
        )
 | 
						|
 | 
						|
        # Assume the new role
 | 
						|
        iam_account_b = boto3.client(
 | 
						|
            "iam",
 | 
						|
            aws_access_key_id=response["Credentials"]["AccessKeyId"],
 | 
						|
            aws_secret_access_key=response["Credentials"]["SecretAccessKey"],
 | 
						|
            aws_session_token=response["Credentials"]["SessionToken"],
 | 
						|
            region_name="us-east-1",
 | 
						|
        )
 | 
						|
 | 
						|
        # Verify new users belong to the different account
 | 
						|
        user = iam_account_b.create_user(UserName="user-in-new-account")["User"]
 | 
						|
        user["Arn"].should.equal(
 | 
						|
            f"arn:aws:iam::{self.account_b}:user/user-in-new-account"
 | 
						|
        )
 | 
						|
 | 
						|
    def test_assume_role_with_saml_in_different_account(self):
 | 
						|
        role_name = "test-role"
 | 
						|
        provider_name = "TestProvFed"
 | 
						|
        fed_identifier = "7ca82df9-1bad-4dd3-9b2b-adb68b554282"
 | 
						|
        fed_name = "testuser"
 | 
						|
        role_input = "arn:aws:iam::{account_id}:role/{role_name}".format(
 | 
						|
            account_id=self.account_b, role_name=role_name
 | 
						|
        )
 | 
						|
        principal_role = (
 | 
						|
            "arn:aws:iam:{account_id}:saml-provider/{provider_name}".format(
 | 
						|
                account_id=ACCOUNT_ID, provider_name=provider_name
 | 
						|
            )
 | 
						|
        )
 | 
						|
        saml_assertion = """<?xml version="1.0"?>
 | 
						|
        <samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" ID="_00000000-0000-0000-0000-000000000000" Version="2.0" IssueInstant="2012-01-01T12:00:00.000Z" Destination="https://signin.aws.amazon.com/saml" Consent="urn:oasis:names:tc:SAML:2.0:consent:unspecified">
 | 
						|
          <Issuer xmlns="urn:oasis:names:tc:SAML:2.0:assertion">http://localhost/</Issuer>
 | 
						|
          <samlp:Status>
 | 
						|
            <samlp:StatusCode Value="urn:oasis:names:tc:SAML:2.0:status:Success"/>
 | 
						|
          </samlp:Status>
 | 
						|
          <Assertion xmlns="urn:oasis:names:tc:SAML:2.0:assertion" ID="_00000000-0000-0000-0000-000000000000" IssueInstant="2012-12-01T12:00:00.000Z" Version="2.0">
 | 
						|
            <Issuer>http://localhost:3000/</Issuer>
 | 
						|
            <ds:Signature xmlns:ds="http://www.w3.org/2000/09/xmldsig#">
 | 
						|
              <ds:SignedInfo>
 | 
						|
                <ds:CanonicalizationMethod Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
 | 
						|
                <ds:SignatureMethod Algorithm="http://www.w3.org/2001/04/xmldsig-more#rsa-sha256"/>
 | 
						|
                <ds:Reference URI="#_00000000-0000-0000-0000-000000000000">
 | 
						|
                  <ds:Transforms>
 | 
						|
                    <ds:Transform Algorithm="http://www.w3.org/2000/09/xmldsig#enveloped-signature"/>
 | 
						|
                    <ds:Transform Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
 | 
						|
                  </ds:Transforms>
 | 
						|
                  <ds:DigestMethod Algorithm="http://www.w3.org/2001/04/xmlenc#sha256"/>
 | 
						|
                  <ds:DigestValue>NTIyMzk0ZGI4MjI0ZjI5ZGNhYjkyOGQyZGQ1NTZjODViZjk5YTY4ODFjOWRjNjkyYzZmODY2ZDQ4NjlkZjY3YSAgLQo=</ds:DigestValue>
 | 
						|
                </ds:Reference>
 | 
						|
              </ds:SignedInfo>
 | 
						|
              <ds:SignatureValue>NTIyMzk0ZGI4MjI0ZjI5ZGNhYjkyOGQyZGQ1NTZjODViZjk5YTY4ODFjOWRjNjkyYzZmODY2ZDQ4NjlkZjY3YSAgLQo=</ds:SignatureValue>
 | 
						|
              <KeyInfo xmlns="http://www.w3.org/2000/09/xmldsig#">
 | 
						|
                <ds:X509Data>
 | 
						|
                  <ds:X509Certificate>NTIyMzk0ZGI4MjI0ZjI5ZGNhYjkyOGQyZGQ1NTZjODViZjk5YTY4ODFjOWRjNjkyYzZmODY2ZDQ4NjlkZjY3YSAgLQo=</ds:X509Certificate>
 | 
						|
                </ds:X509Data>
 | 
						|
              </KeyInfo>
 | 
						|
            </ds:Signature>
 | 
						|
            <Subject>
 | 
						|
              <NameID Format="urn:oasis:names:tc:SAML:2.0:nameid-format:persistent">{fed_identifier}</NameID>
 | 
						|
              <SubjectConfirmation Method="urn:oasis:names:tc:SAML:2.0:cm:bearer">
 | 
						|
                <SubjectConfirmationData NotOnOrAfter="2012-01-01T13:00:00.000Z" Recipient="https://signin.aws.amazon.com/saml"/>
 | 
						|
              </SubjectConfirmation>
 | 
						|
            </Subject>
 | 
						|
            <Conditions NotBefore="2012-01-01T12:00:00.000Z" NotOnOrAfter="2012-01-01T13:00:00.000Z">
 | 
						|
              <AudienceRestriction>
 | 
						|
                <Audience>urn:amazon:webservices</Audience>
 | 
						|
              </AudienceRestriction>
 | 
						|
            </Conditions>
 | 
						|
            <AttributeStatement>
 | 
						|
              <Attribute Name="https://aws.amazon.com/SAML/Attributes/RoleSessionName">
 | 
						|
                <AttributeValue>{fed_name}</AttributeValue>
 | 
						|
              </Attribute>
 | 
						|
              <Attribute Name="https://aws.amazon.com/SAML/Attributes/Role">
 | 
						|
                <AttributeValue>arn:aws:iam::{account_id}:role/{role_name},arn:aws:iam::{account_id}:saml-provider/{provider_name}</AttributeValue>
 | 
						|
              </Attribute>
 | 
						|
              <Attribute Name="https://aws.amazon.com/SAML/Attributes/SessionDuration">
 | 
						|
                <AttributeValue>900</AttributeValue>
 | 
						|
              </Attribute>
 | 
						|
            </AttributeStatement>
 | 
						|
            <AuthnStatement AuthnInstant="2012-01-01T12:00:00.000Z" SessionIndex="_00000000-0000-0000-0000-000000000000">
 | 
						|
              <AuthnContext>
 | 
						|
                <AuthnContextClassRef>urn:oasis:names:tc:SAML:2.0:ac:classes:PasswordProtectedTransport</AuthnContextClassRef>
 | 
						|
              </AuthnContext>
 | 
						|
            </AuthnStatement>
 | 
						|
          </Assertion>
 | 
						|
        </samlp:Response>""".format(
 | 
						|
            account_id=self.account_b,
 | 
						|
            role_name=role_name,
 | 
						|
            provider_name=provider_name,
 | 
						|
            fed_identifier=fed_identifier,
 | 
						|
            fed_name=fed_name,
 | 
						|
        ).replace(
 | 
						|
            "\n", ""
 | 
						|
        )
 | 
						|
 | 
						|
        assume_role_response = self.sts.assume_role_with_saml(
 | 
						|
            RoleArn=role_input,
 | 
						|
            PrincipalArn=principal_role,
 | 
						|
            SAMLAssertion=b64encode(saml_assertion.encode("utf-8")).decode("utf-8"),
 | 
						|
        )
 | 
						|
 | 
						|
        # Assume the new role
 | 
						|
        iam_account_b = boto3.client(
 | 
						|
            "iam",
 | 
						|
            aws_access_key_id=assume_role_response["Credentials"]["AccessKeyId"],
 | 
						|
            aws_secret_access_key=assume_role_response["Credentials"][
 | 
						|
                "SecretAccessKey"
 | 
						|
            ],
 | 
						|
            aws_session_token=assume_role_response["Credentials"]["SessionToken"],
 | 
						|
            region_name="us-east-1",
 | 
						|
        )
 | 
						|
 | 
						|
        # Verify new users belong to the different account
 | 
						|
        user = iam_account_b.create_user(UserName="user-in-new-account")["User"]
 | 
						|
        user["Arn"].should.equal(
 | 
						|
            f"arn:aws:iam::{self.account_b}:user/user-in-new-account"
 | 
						|
        )
 | 
						|
 | 
						|
    def test_dynamodb_supports_multiple_accounts(self):
 | 
						|
        ddb_client = boto3.client("dynamodb", region_name="us-east-1")
 | 
						|
 | 
						|
        ddb_client.create_table(
 | 
						|
            TableName="table-in-default-account",
 | 
						|
            KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}],
 | 
						|
            AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}],
 | 
						|
            ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 5},
 | 
						|
        )
 | 
						|
        # assume role to another aws account
 | 
						|
        role_name = f"arn:aws:iam::{self.account_b}:role/my-role"
 | 
						|
        response = self.sts.assume_role(
 | 
						|
            RoleArn=role_name,
 | 
						|
            RoleSessionName="test-session-name",
 | 
						|
            ExternalId="test-external-id",
 | 
						|
        )
 | 
						|
 | 
						|
        # Assume the new role
 | 
						|
        ddb_account_b = boto3.client(
 | 
						|
            "dynamodb",
 | 
						|
            aws_access_key_id=response["Credentials"]["AccessKeyId"],
 | 
						|
            aws_secret_access_key=response["Credentials"]["SecretAccessKey"],
 | 
						|
            aws_session_token=response["Credentials"]["SessionToken"],
 | 
						|
            region_name="us-east-1",
 | 
						|
        )
 | 
						|
 | 
						|
        # Verify new dynamodb belong to the different account
 | 
						|
        ddb_account_b.create_table(
 | 
						|
            TableName="table-in-new-account",
 | 
						|
            KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}],
 | 
						|
            AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}],
 | 
						|
            ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 5},
 | 
						|
        )
 | 
						|
 | 
						|
        table = ddb_client.describe_table(TableName="table-in-default-account")["Table"]
 | 
						|
        table["TableArn"].should.equal(
 | 
						|
            "arn:aws:dynamodb:us-east-1:123456789012:table/table-in-default-account"
 | 
						|
        )
 | 
						|
 | 
						|
        table = ddb_account_b.describe_table(TableName="table-in-new-account")["Table"]
 | 
						|
        table["TableArn"].should.equal(
 | 
						|
            f"arn:aws:dynamodb:us-east-1:{self.account_b}:table/table-in-new-account"
 | 
						|
        )
 |