Use STS to fetch the account ID instead of IAM

philchristensen · web-flow · commit f45927c79119 · 2016-12-06T12:02:03.000-05:00
This approach will work with both normal IAM account credentials and assumed roles.
diff --git a/aws_lambda/aws_lambda.py b/aws_lambda/aws_lambda.py
@@ -264,9 +264,9 @@ def get_role_name(account_id, role):
 
 
 def get_account_id(aws_access_key_id, aws_secret_access_key):
-    """Query IAM for a users' account_id"""
-    client = get_client('iam', aws_access_key_id, aws_secret_access_key)
-    return client.get_user()['User']['Arn'].split(':')[4]
+    """Query STS for a users' account_id"""
+    client = get_client('sts', aws_access_key_id, aws_secret_access_key)
+    return client.get_caller_identity().get('Account')
 
 
 def get_client(client, aws_access_key_id, aws_secret_access_key, region=None):
