From d9cb1f2d353ccc8d0ef77f5a98fc2d8f4ca82252 Mon Sep 17 00:00:00 2001
From: acsbendi <acsbendi28@gmail.com>
Date: Wed, 21 Aug 2019 10:45:36 +0200
Subject: [PATCH] Implemented returning random assumed role ID.

---
 moto/sts/models.py         |  3 ++-
 moto/sts/responses.py      |  2 +-
 moto/sts/utils.py          | 20 +++++++++++++++-----
 tests/test_sts/test_sts.py | 30 +++++++++++++++++-------------
 4 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/moto/sts/models.py b/moto/sts/models.py
index e437575b..ff9de84b 100644
--- a/moto/sts/models.py
+++ b/moto/sts/models.py
@@ -2,7 +2,7 @@ from __future__ import unicode_literals
 import datetime
 from moto.core import BaseBackend, BaseModel
 from moto.core.utils import iso_8601_datetime_with_milliseconds
-from moto.sts.utils import random_access_key_id, random_secret_access_key, random_session_token
+from moto.sts.utils import random_access_key_id, random_secret_access_key, random_session_token, random_assumed_role_id
 
 
 class Token(BaseModel):
@@ -30,6 +30,7 @@ class AssumedRole(BaseModel):
         self.access_key_id = "ASIA" + random_access_key_id()
         self.secret_access_key = random_secret_access_key()
         self.session_token = random_session_token()
+        self.assumed_role_id = "AROA" + random_assumed_role_id()
 
     @property
     def expiration_ISO8601(self):
diff --git a/moto/sts/responses.py b/moto/sts/responses.py
index 24ec181b..c04d1636 100644
--- a/moto/sts/responses.py
+++ b/moto/sts/responses.py
@@ -91,7 +91,7 @@ ASSUME_ROLE_RESPONSE = """<AssumeRoleResponse xmlns="https://sts.amazonaws.com/d
     </Credentials>
     <AssumedRoleUser>
       <Arn>{{ role.arn }}</Arn>
-      <AssumedRoleId>ARO123EXAMPLE123:{{ role.session_name }}</AssumedRoleId>
+      <AssumedRoleId>{{ role.assumed_role_id }}:{{ role.session_name }}</AssumedRoleId>
     </AssumedRoleUser>
     <PackedPolicySize>6</PackedPolicySize>
   </AssumeRoleResult>
diff --git a/moto/sts/utils.py b/moto/sts/utils.py
index 8e612972..50767729 100644
--- a/moto/sts/utils.py
+++ b/moto/sts/utils.py
@@ -6,15 +6,12 @@ import string
 import six
 
 ACCOUNT_SPECIFIC_ACCESS_KEY_PREFIX = "8NWMTLYQ"
+ACCOUNT_SPECIFIC_ASSUMED_ROLE_ID_PREFIX = "3X42LBCD"
 SESSION_TOKEN_PREFIX = "FQoGZXIvYXdzEBYaD"
 
 
 def random_access_key_id():
-    return ACCOUNT_SPECIFIC_ACCESS_KEY_PREFIX + ''.join(six.text_type(
-        random.choice(
-            string.ascii_uppercase + string.digits
-        )) for _ in range(8)
-    )
+    return ACCOUNT_SPECIFIC_ACCESS_KEY_PREFIX + _random_uppercase_or_digit_sequence(8)
 
 
 def random_secret_access_key():
@@ -23,3 +20,16 @@ def random_secret_access_key():
 
 def random_session_token():
     return SESSION_TOKEN_PREFIX + base64.b64encode(os.urandom(266))[len(SESSION_TOKEN_PREFIX):].decode()
+
+
+def random_assumed_role_id():
+    return ACCOUNT_SPECIFIC_ASSUMED_ROLE_ID_PREFIX + _random_uppercase_or_digit_sequence(9)
+
+
+def _random_uppercase_or_digit_sequence(length):
+    return ''.join(
+        six.text_type(
+            random.choice(
+                string.ascii_uppercase + string.digits
+            )) for _ in range(length)
+    )
diff --git a/tests/test_sts/test_sts.py b/tests/test_sts/test_sts.py
index 566f17ff..3cc44b99 100644
--- a/tests/test_sts/test_sts.py
+++ b/tests/test_sts/test_sts.py
@@ -40,10 +40,12 @@ def test_get_federation_token():
 
 
 @freeze_time("2012-01-01 12:00:00")
-@mock_sts_deprecated
+@mock_sts
 def test_assume_role():
-    conn = boto.connect_sts()
+    client = boto3.client(
+        "sts", region_name='us-east-1')
 
+    session_name = "session-name"
     policy = json.dumps({
         "Statement": [
             {
@@ -59,19 +61,21 @@ def test_assume_role():
         ]
     })
     s3_role = "arn:aws:iam::123456789012:role/test-role"
-    role = conn.assume_role(s3_role, "session-name",
-                            policy, duration_seconds=123)
+    assume_role_response = client.assume_role(RoleArn=s3_role, RoleSessionName=session_name,
+                                              Policy=policy, DurationSeconds=900)
 
-    credentials = role.credentials
-    credentials.expiration.should.equal('2012-01-01T12:02:03.000Z')
-    credentials.session_token.should.have.length_of(356)
-    assert credentials.session_token.startswith("FQoGZXIvYXdzE")
-    credentials.access_key.should.have.length_of(20)
-    assert credentials.access_key.startswith("ASIA")
-    credentials.secret_key.should.have.length_of(40)
+    credentials = assume_role_response['Credentials']
+    credentials['Expiration'].isoformat().should.equal('2012-01-01T12:15:00+00:00')
+    credentials['SessionToken'].should.have.length_of(356)
+    assert credentials['SessionToken'].startswith("FQoGZXIvYXdzE")
+    credentials['AccessKeyId'].should.have.length_of(20)
+    assert credentials['AccessKeyId'].startswith("ASIA")
+    credentials['SecretAccessKey'].should.have.length_of(40)
 
-    role.user.arn.should.equal("arn:aws:iam::123456789012:role/test-role")
-    role.user.assume_role_id.should.contain("session-name")
+    assume_role_response['AssumedRoleUser']['Arn'].should.equal("arn:aws:iam::123456789012:role/test-role")
+    assert assume_role_response['AssumedRoleUser']['AssumedRoleId'].startswith("AROA")
+    assert assume_role_response['AssumedRoleUser']['AssumedRoleId'].endswith(":" + session_name)
+    assume_role_response['AssumedRoleUser']['AssumedRoleId'].should.have.length_of(21 + 1 + len(session_name))
 
 
 @mock_sts