From 13f5e930a44304a5209ec24d7c4d7aca6513ff80 Mon Sep 17 00:00:00 2001
From: rfrench42 <31379169+rfrench42@users.noreply.github.com>
Date: Mon, 19 Aug 2019 15:12:31 -0400
Subject: [PATCH] Added support for addtional partitions in ARN (#28699)
 (#28704)

---
 lib/ansible/modules/cloud/amazon/lambda.py | 29 ++++++++++++++--------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/lib/ansible/modules/cloud/amazon/lambda.py b/lib/ansible/modules/cloud/amazon/lambda.py
index 3264296ba5f..7fed40df69e 100644
--- a/lib/ansible/modules/cloud/amazon/lambda.py
+++ b/lib/ansible/modules/cloud/amazon/lambda.py
@@ -206,6 +206,7 @@ from ansible.module_utils.ec2 import compare_aws_tags
 import base64
 import hashlib
 import traceback
+import re
 
 try:
     from botocore.exceptions import ClientError, BotoCoreError, ValidationError, ParamValidationError
@@ -213,33 +214,41 @@ except ImportError:
     pass  # protected by AnsibleAWSModule
 
 
-def get_account_id(module, region=None, endpoint=None, **aws_connect_kwargs):
-    """return the account id we are currently working on
+def get_account_info(module, region=None, endpoint=None, **aws_connect_kwargs):
+    """return the account information (account id and partition) we are currently working on
 
-    get_account_id tries too find out the account that we are working
+    get_account_info tries too find out the account that we are working
     on.  It's not guaranteed that this will be easy so we try in
     several different ways.  Giving either IAM or STS privileges to
     the account should be enough to permit this.
     """
     account_id = None
+    partition = None
     try:
         sts_client = boto3_conn(module, conn_type='client', resource='sts',
                                 region=region, endpoint=endpoint, **aws_connect_kwargs)
-        account_id = sts_client.get_caller_identity().get('Account')
+        caller_id = sts_client.get_caller_identity()
+        account_id = caller_id.get('Account')
+        partition = caller_id.get('Arn').split(':')[1]
     except ClientError:
         try:
             iam_client = boto3_conn(module, conn_type='client', resource='iam',
                                     region=region, endpoint=endpoint, **aws_connect_kwargs)
-            account_id = iam_client.get_user()['User']['Arn'].split(':')[4]
+            arn, partition, service, reg, account_id, resource = iam_client.get_user()['User']['Arn'].split(':')
         except ClientError as e:
             if (e.response['Error']['Code'] == 'AccessDenied'):
                 except_msg = to_native(e.message)
-                account_id = except_msg.search(r"arn:aws:iam::([0-9]{12,32}):\w+/").group(1)
+                m = except_msg.search(r"arn:(aws(-([a-z\-]+))?):iam::([0-9]{12,32}):\w+/")
+                account_id = m.group(4)
+                partition = m.group(1)
             if account_id is None:
                 module.fail_json_aws(e, msg="getting account information")
+            if partition is None:
+                module.fail_json_aws(e, msg="getting account information: partition")
         except Exception as e:
             module.fail_json_aws(e, msg="getting account information")
-    return account_id
+
+    return account_id, partition
 
 
 def get_current_function(connection, function_name, qualifier=None):
@@ -377,12 +386,12 @@ def main():
         module.fail_json_aws(e, msg="Trying to connect to AWS")
 
     if state == 'present':
-        if role.startswith('arn:aws:iam'):
+        if re.match(r'^arn:aws(-([a-z\-]+))?:iam', role):
             role_arn = role
         else:
             # get account ID and assemble ARN
-            account_id = get_account_id(module, region=region, endpoint=ec2_url, **aws_connect_kwargs)
-            role_arn = 'arn:aws:iam::{0}:role/{1}'.format(account_id, role)
+            account_id, partition = get_account_info(module, region=region, endpoint=ec2_url, **aws_connect_kwargs)
+            role_arn = 'arn:{0}:iam::{1}:role/{2}'.format(partition, account_id, role)
 
     # Get function configuration if present, False otherwise
     current_function = get_current_function(client, name)