chrisoverzero · chrisoverzero · Apr 13, 2021
diff --git a/docs/internals/generated_resources.rst b/docs/internals/generated_resources.rst
@@ -86,6 +86,31 @@ AWS::IAM::Role                     CodeDeployServiceRole
 
   NOTE: ``AWS::IAM::Role`` resources are only generated if no Role parameter is supplied for DeploymentPreference
 
+
+With KmsKeyArn Property
+~~~~~~~~~~~~~~~~~~~~~~~
+
+Example:
+
+.. code:: yaml
+
+  MyFunction:
+    Type: AWS::Serverless::Function
+    Properties:
+      ...
+      KmsKeyArn: !GetAtt MyKmsKey.Arn
+      ...
+
+
+Additional generated resources:
+
+================================== ================================
+CloudFormation Resource Type       Logical ID 
+================================== ================================
+AWS::IAM::Poilicy                  MyFunction\ **DecryptEnvironmentVariablesPolicy**
+================================== ================================
+
+
 With Events
 ~~~~~~~~~~~
 

diff --git a/samtranslator/model/iam.py b/samtranslator/model/iam.py
@@ -17,6 +17,17 @@ class IAMRole(Resource):
     runtime_attrs = {"name": lambda self: ref(self.logical_id), "arn": lambda self: fnGetAtt(self.logical_id, "Arn")}
 
 
+class IAMPolicy(Resource):
+    resource_type = "AWS::IAM::Policy"
+    property_types = {
+        "PolicyDocument": PropertyType(True, is_type(object)),
+        "PolicyName": PropertyType(True, is_str()),
+        "Roles": PropertyType(True, is_type(list)),
+    }
+
+    runtime_attrs = {"name": lambda self: ref(self.logical_id)}
+
+
 class IAMRolePolicies:
     @classmethod
     def construct_assume_role_policy_for_service_principal(cls, service_principal):
@@ -129,3 +140,27 @@ def lambda_invoke_function_role_policy(cls, function_arn, logical_id):
             },
         }
         return document
+
+
+class IAMPolicies:
+    @classmethod
+    def lambda_decrypt_environment_variables_policy(cls, logical_id, kms_key_arn, function_arn, role_arn):
+        policy_name = "DecryptEnvironmentVariablesPolicy"
+        policy = IAMPolicy(logical_id + policy_name)
+        policy.PolicyName = policy_name
+        policy.PolicyDocument = {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Action": "kms:Decrypt",
+                    "Effect": "Allow",
+                    "Resource": kms_key_arn,
+                    "Condition": {
+                        "StringEquals": {"kms:ViaService": "lambda.amazonaws.com"},
+                        "ForAnyValue:ArnEquals": {"kms:EncryptionContext:aws:lambda:FunctionArn": function_arn},
+                    },
+                }
+            ],
+        }
+        policy.Roles = [role_arn]
+        return policy
diff --git a/samtranslator/model/sam_resources.py b/samtranslator/model/sam_resources.py
@@ -24,7 +24,7 @@
 from samtranslator.model.dynamodb import DynamoDBTable
 from samtranslator.model.exceptions import InvalidEventException, InvalidResourceException
 from samtranslator.model.resource_policies import ResourcePolicies, PolicyTypes
-from samtranslator.model.iam import IAMRole, IAMRolePolicies
+from samtranslator.model.iam import IAMRole, IAMPolicy, IAMRolePolicies, IAMPolicies
 from samtranslator.model.lambda_ import (
     LambdaFunction,
     LambdaVersion,
@@ -183,6 +183,10 @@ def to_cloudformation(self, **kwargs):
             lambda_function.Role = execution_role.get_runtime_attr("arn")
             resources.append(execution_role)
 
+        if self.KmsKeyArn:
+            env_vars_policy = self._construct_environment_variables_decrypt_policy(lambda_function)
+            resources.append(env_vars_policy)
+
         try:
             resources += self._generate_event_resources(
                 lambda_function,
@@ -493,6 +497,18 @@ def _construct_role(self, managed_policy_map, event_invoke_policies):
         )
         return execution_role
 
+    def _construct_environment_variables_decrypt_policy(self, lambda_function):
+        """Constructs a Lambda policy for reading encrypted environment variables
+        based on this SAM function's KmsKeyArn property.
+
+        :returns: the generated IAM Policy
+        :rtype: model.iam.IAMPolicy
+        """
+        policy = IAMPolicies.lambda_decrypt_environment_variables_policy(
+            self.logical_id, lambda_function.KmsKeyArn, lambda_function.get_runtime_attr("arn"), lambda_function.Role
+        )
+        return policy
+
     def _validate_package_type(self, lambda_function):
         """
         Validates Function based on the existence of Package type

diff --git a/tests/translator/output/aws-cn/function_with_kmskeyarn.json b/tests/translator/output/aws-cn/function_with_kmskeyarn.json
@@ -1,17 +1,55 @@
 {
   "Resources": {
-    "FunctionWithKeyArnRole": {
-      "Type": "AWS::IAM::Role",
+    "myKey": {
+      "Type": "AWS::KMS::Key",
       "Properties": {
-        "ManagedPolicyArns": [
-          "arn:aws-cn:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
-        ],
+        "Description": "A sample key",
+        "KeyPolicy": {
+          "Version": "2012-10-17",
+          "Id": "key-default-1",
+          "Statement": [
+            {
+              "Sid": "Allow administration of the key",
+              "Effect": "Allow",
+              "Principal": {
+                "AWS": "arn:aws:iam::123456789012:user/Alice"
+              },
+              "Action": [
+                "kms:Create*"
+              ],
+              "Resource": "*"
+            }
+          ]
+        }
+      }
+    },
+    "FunctionWithKeyArn": {
+      "Type": "AWS::Lambda::Function",
+      "Properties": {
+        "Code": {
+          "S3Bucket": "sam-demo-bucket",
+          "S3Key": "hello.zip"
+        },
+        "Handler": "hello.handler",
+        "Role": {
+          "Fn::GetAtt": [
+            "FunctionWithKeyArnRole",
+            "Arn"
+          ]
+        },
+        "Runtime": "python2.7",
         "Tags": [
           {
-            "Value": "SAM",
-            "Key": "lambda:createdBy"
+            "Key": "lambda:createdBy",
+            "Value": "SAM"
           }
         ],
+        "KmsKeyArn": "thisIsaKey"
+      }
+    },
+    "FunctionWithKeyArnRole": {
+      "Type": "AWS::IAM::Role",
+      "Properties": {
         "AssumeRolePolicyDocument": {
           "Version": "2012-10-17",
           "Statement": [
@@ -27,111 +65,149 @@
               }
             }
           ]
-        }
-      }
-    },
-    "FunctionWithReferenceToKeyArnRole": {
-      "Type": "AWS::IAM::Role",
-      "Properties": {
+        },
         "ManagedPolicyArns": [
           "arn:aws-cn:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
         ],
         "Tags": [
           {
-            "Value": "SAM",
-            "Key": "lambda:createdBy"
+            "Key": "lambda:createdBy",
+            "Value": "SAM"
           }
-        ],
-        "AssumeRolePolicyDocument": {
+        ]
+      }
+    },
+    "FunctionWithKeyArnDecryptEnvironmentVariablesPolicy": {
+      "Type": "AWS::IAM::Policy",
+      "Properties": {
+        "PolicyDocument": {
           "Version": "2012-10-17",
           "Statement": [
             {
-              "Action": [
-                "sts:AssumeRole"
-              ],
+              "Action": "kms:Decrypt",
               "Effect": "Allow",
-              "Principal": {
-                "Service": [
-                  "lambda.amazonaws.com"
-                ]
+              "Resource": "thisIsaKey",
+              "Condition": {
+                "StringEquals": {
+                  "kms:ViaService": "lambda.amazonaws.com"
+                },
+                "ForAnyValue:ArnEquals": {
+                  "kms:EncryptionContext:aws:lambda:FunctionArn": {
+                    "Fn::GetAtt": [
+                      "FunctionWithKeyArn",
+                      "Arn"
+                    ]
+                  }
+                }
               }
             }
           ]
-        }
+        },
+        "PolicyName": "DecryptEnvironmentVariablesPolicy",
+        "Roles": [
+          {
+            "Fn::GetAtt": [
+              "FunctionWithKeyArnRole",
+              "Arn"
+            ]
+          }
+        ]
       }
     },
-    "FunctionWithKeyArn": {
+    "FunctionWithReferenceToKeyArn": {
       "Type": "AWS::Lambda::Function",
       "Properties": {
         "Code": {
           "S3Bucket": "sam-demo-bucket",
           "S3Key": "hello.zip"
         },
-        "Tags": [
-          {
-            "Value": "SAM",
-            "Key": "lambda:createdBy"
-          }
-        ],
-        "KmsKeyArn": "thisIsaKey",
         "Handler": "hello.handler",
         "Role": {
           "Fn::GetAtt": [
-            "FunctionWithKeyArnRole",
+            "FunctionWithReferenceToKeyArnRole",
             "Arn"
           ]
         },
-        "Runtime": "python2.7"
+        "Runtime": "python2.7",
+        "Tags": [
+          {
+            "Key": "lambda:createdBy",
+            "Value": "SAM"
+          }
+        ],
+        "KmsKeyArn": {
+          "Ref": "myKey"
+        }
       }
     },
-    "myKey": {
-      "Type": "AWS::KMS::Key",
+    "FunctionWithReferenceToKeyArnRole": {
+      "Type": "AWS::IAM::Role",
       "Properties": {
-        "KeyPolicy": {
+        "AssumeRolePolicyDocument": {
           "Version": "2012-10-17",
-          "Id": "key-default-1",
           "Statement": [
             {
               "Action": [
-                "kms:Create*"
+                "sts:AssumeRole"
               ],
-              "Sid": "Allow administration of the key",
-              "Resource": "*",
               "Effect": "Allow",
               "Principal": {
-                "AWS": "arn:aws:iam::123456789012:user/Alice"
+                "Service": [
+                  "lambda.amazonaws.com"
+                ]
               }
             }
           ]
         },
-        "Description": "A sample key"
+        "ManagedPolicyArns": [
+          "arn:aws-cn:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
+        ],
+        "Tags": [
+          {
+            "Key": "lambda:createdBy",
+            "Value": "SAM"
+          }
+        ]
       }
     },
-    "FunctionWithReferenceToKeyArn": {
-      "Type": "AWS::Lambda::Function",
+    "FunctionWithReferenceToKeyArnDecryptEnvironmentVariablesPolicy": {
+      "Type": "AWS::IAM::Policy",
       "Properties": {
-        "Code": {
-          "S3Bucket": "sam-demo-bucket",
-          "S3Key": "hello.zip"
+        "PolicyDocument": {
+          "Version": "2012-10-17",
+          "Statement": [
+            {
+              "Action": "kms:Decrypt",
+              "Effect": "Allow",
+              "Resource": {
+                "Ref": "myKey"
+              },
+              "Condition": {
+                "StringEquals": {
+                  "kms:ViaService": "lambda.amazonaws.com"
+                },
+                "ForAnyValue:ArnEquals": {
+                  "kms:EncryptionContext:aws:lambda:FunctionArn": {
+                    "Fn::GetAtt": [
+                      "FunctionWithReferenceToKeyArn",
+                      "Arn"
+                    ]
+                  }
+                }
+              }
+            }
+          ]
         },
-        "Tags": [
+        "PolicyName": "DecryptEnvironmentVariablesPolicy",
+        "Roles": [
           {
-            "Value": "SAM",
-            "Key": "lambda:createdBy"
+            "Fn::GetAtt": [
+              "FunctionWithReferenceToKeyArnRole",
+              "Arn"
+            ]
           }
-        ],
-        "KmsKeyArn": {
-          "Ref": "myKey"
-        },
-        "Handler": "hello.handler",
-        "Role": {
-          "Fn::GetAtt": [
-            "FunctionWithReferenceToKeyArnRole",
-            "Arn"
-          ]
-        },
-        "Runtime": "python2.7"
+        ]
       }
     }
   }
-}
+}