Merge branch 'master' of github.com:nficano/python-lambda

nficano · nficano · commit 9ba608c212f1 · 2017-09-18T11:57:56.000-04:00
* 'master' of github.com:nficano/python-lambda:
  modify the role ARN slightly if in govcloud. New GovCloud regions may be added, so included a constant map ARN_PREFIXES for future expansion
diff --git a/aws_lambda/aws_lambda.py b/aws_lambda/aws_lambda.py
@@ -24,6 +24,10 @@
 from .helpers import timestamp
 
 
+ARN_PREFIXES = {
+    'us-gov-west-1': 'aws-us-gov'
+}
+
 log = logging.getLogger(__name__)
 
 
@@ -358,9 +362,10 @@ def pip_install_to_target(path, requirements=False, local_package=None):
     _install_packages(path, packages)
 
 
-def get_role_name(account_id, role):
+def get_role_name(region, account_id, role):
     """Shortcut to insert the `account_id` and `role` into the iam string."""
-    return 'arn:aws:iam::{0}:role/{1}'.format(account_id, role)
+    prefix = ARN_PREFIXES.get(region, 'aws')
+    return 'arn:{0}:iam::{1}:role/{2}'.format(prefix, account_id, role)
 
 
 def get_account_id(aws_access_key_id, aws_secret_access_key):
@@ -389,7 +394,7 @@ def create_function(cfg, path_to_zip_file):
     aws_secret_access_key = cfg.get('aws_secret_access_key')
 
     account_id = get_account_id(aws_access_key_id, aws_secret_access_key)
-    role = get_role_name(account_id, cfg.get('role', 'lambda_basic_execution'))
+    role = get_role_name(cfg.get('region'), account_id, cfg.get('role', 'lambda_basic_execution'))
 
     client = get_client(
         'lambda', aws_access_key_id, aws_secret_access_key,
@@ -436,7 +441,7 @@ def update_function(cfg, path_to_zip_file):
     aws_secret_access_key = cfg.get('aws_secret_access_key')
 
     account_id = get_account_id(aws_access_key_id, aws_secret_access_key)
-    role = get_role_name(account_id, cfg.get('role', 'lambda_basic_execution'))
+    role = get_role_name(cfg.get('region'), account_id, cfg.get('role', 'lambda_basic_execution'))
 
     client = get_client(
         'lambda', aws_access_key_id, aws_secret_access_key,
