Update planner to include source account for S3 events

Author: Jon Beilke

chalice/awsclient.py

@@ -1522,24 +1522,27 @@ def _merge_s3_notification_config(self, existing_config, new_config):
             final_config.append(new_config)
         return final_config
 
-    def add_permission_for_s3_event(self, bucket, function_arn):
-        # type: (str, str) -> None
+    def add_permission_for_s3_event(self, bucket, function_arn, account_id):
+        # type: (str, str, str) -> None
         bucket_arn = 'arn:{partition}:s3:::{bucket}'.format(
             partition=self.partition_name, bucket=bucket)
         self._add_lambda_permission_if_needed(
             source_arn=bucket_arn,
             function_arn=function_arn,
             service_name='s3',
+            source_account=account_id,
         )
 
-    def remove_permission_for_s3_event(self, bucket, function_arn):
-        # type: (str, str) -> None
+    def remove_permission_for_s3_event(self, bucket, function_arn,
+                                       account_id):
+        # type: (str, str, str) -> None
         bucket_arn = 'arn:{partition}:s3:::{bucket}'.format(
             partition=self.partition_name, bucket=bucket)
         self._remove_lambda_permission_if_needed(
             source_arn=bucket_arn,
             function_arn=function_arn,
             service_name='s3',
+            source_account=account_id,
         )
 
     def disconnect_s3_bucket_from_lambda(self, bucket, function_arn):
@@ -1562,22 +1565,25 @@ def disconnect_s3_bucket_from_lambda(self, bucket, function_arn):
         )
 
     def _add_lambda_permission_if_needed(self, source_arn, function_arn,
-                                         service_name):
-        # type: (str, str, str) -> None
+                                         service_name, source_account=None):
+        # type: (str, str, str, Optional[str]) -> None
         policy = self.get_function_policy(function_arn)
         if self._policy_gives_access(policy, source_arn, service_name):
             return
         random_id = self._random_id()
         dns_suffix = self.endpoint_dns_suffix_from_arn(source_arn)
-        self._client('lambda').add_permission(
-            Action='lambda:InvokeFunction',
-            FunctionName=function_arn,
-            StatementId=random_id,
-            Principal=self.service_principal(service_name,
-                                             self.region_name,
-                                             dns_suffix),
-            SourceArn=source_arn,
-        )
+        kwargs = {
+            'Action': 'lambda:InvokeFunction',
+            'FunctionName': function_arn,
+            'StatementId': random_id,
+            'Principal': self.service_principal(service_name,
+                                                self.region_name,
+                                                dns_suffix),
+            'SourceArn': source_arn,
+        }
+        if source_account is not None:
+            kwargs['SourceAccount'] = source_account
+        self._client('lambda').add_permission(**kwargs)
 
     def _policy_gives_access(self, policy, source_arn, service_name):
         # type: (Dict[str, Any], str, str) -> bool
@@ -1610,8 +1616,9 @@ def _policy_gives_access(self, policy, source_arn, service_name):
                 return True
         return False
 
-    def _statement_gives_arn_access(self, statement, source_arn, service_name):
-        # type: (Dict[str, Any], str, str) -> bool
+    def _statement_gives_arn_access(self, statement, source_arn,
+                                    service_name, source_account=None):
+        # type: (Dict[str, Any], str, str, Optional[str]) -> bool
         dns_suffix = self.endpoint_dns_suffix_from_arn(source_arn)
         principal = self.service_principal(service_name,
                                            self.region_name,
@@ -1623,19 +1630,29 @@ def _statement_gives_arn_access(self, statement, source_arn, service_name):
             return False
         if statement.get('Principal', {}).get('Service', '') != principal:
             return False
+        if source_account is not None:
+            if statement.get('Condition', {}).get('StringEquals', {}).get(
+                    'AWS:SourceAccount', '') != source_account:
+                return False
         # We're not checking the "Resource" key because we're assuming
         # that lambda.get_policy() is returning the policy for the particular
         # resource in question.
         return True
 
     def _remove_lambda_permission_if_needed(self, source_arn, function_arn,
-                                            service_name):
-        # type: (str, str, str) -> None
+                                            service_name, source_account=None):
+        # type: (str, str, str, Optional[str]) -> None
         client = self._client('lambda')
         policy = self.get_function_policy(function_arn)
         for statement in policy.get('Statement', []):
-            if self._statement_gives_arn_access(statement, source_arn,
-                                                service_name):
+            kwargs = {
+                'statement': statement,
+                'source_arn': source_arn,
+                'service_name': service_name,
+            }
+            if source_account is not None:
+                kwargs['source_account'] = source_account
+            if self._statement_gives_arn_access(**kwargs):
                 client.remove_permission(
                     FunctionName=function_arn,
                     StatementId=statement['Sid'],

chalice/deploy/planner.py

@@ -877,7 +877,8 @@ def _plan_s3bucketnotification(self, resource):
             models.APICall(
                 method_name='add_permission_for_s3_event',
                 params={'bucket': resource.bucket,
-                        'function_arn': function_arn},
+                        'function_arn': function_arn,
+                        'account_id': Variable('account_id')},
             ),
             (models.APICall(
                 method_name='connect_s3_bucket_to_lambda',

tests/functional/test_awsclient.py

@@ -19,8 +19,9 @@
 from chalice.awsclient import ReadTimeout
 
 
-def create_policy_statement(source_arn, service_name, statement_id):
-    return {
+def create_policy_statement(source_arn, service_name, statement_id,
+                            account_id=None):
+    policy_statement = {
         'Action': 'lambda:InvokeFunction',
         'Condition': {
             'ArnLike': {
@@ -32,6 +33,11 @@ def create_policy_statement(source_arn, service_name, statement_id):
         'Resource': 'function-arn',
         'Sid': statement_id,
     }
+    if account_id is not None:
+        policy_statement['Condition']['StringEquals'] = {
+            'AWS:SourceAccount': account_id,
+        }
+    return policy_statement
 
 
 def test_region_name_is_exposed(stubbed_session):
@@ -3256,12 +3262,13 @@ def test_add_permission_for_s3_event(stubbed_session):
         StatementId=stub.ANY,
         Principal='s3.amazonaws.com',
         SourceArn='arn:aws:s3:::mybucket',
+        SourceAccount='12345',
     ).returns({})
     stubbed_session.activate_stubs()
 
     awsclient = TypedAWSClient(stubbed_session)
     awsclient.add_permission_for_s3_event(
-        'mybucket', 'function-arn')
+        'mybucket', 'function-arn', '12345')
     stubbed_session.verify_stubs()
 
 
@@ -3272,6 +3279,9 @@ def test_skip_if_permission_already_granted_to_s3(stubbed_session):
         'Statement': [{
             'Action': 'lambda:InvokeFunction',
             'Condition': {
+                'StringEquals': {
+                    'AWS:SourceAccount': '12345',
+                },
                 'ArnLike': {
                     'AWS:SourceArn': 'arn:aws:s3:::mybucket',
                 }
@@ -3288,7 +3298,7 @@ def test_skip_if_permission_already_granted_to_s3(stubbed_session):
     stubbed_session.activate_stubs()
     awsclient = TypedAWSClient(stubbed_session)
     awsclient.add_permission_for_s3_event(
-        'mybucket', 'function-arn')
+        'mybucket', 'function-arn', '12345')
     stubbed_session.verify_stubs()
 
 
@@ -3480,7 +3490,8 @@ def test_can_remove_s3_permission(stubbed_session):
         'Id': 'default',
         'Statement': [create_policy_statement('arn:aws:s3:::mybucket',
                                               service_name='s3',
-                                              statement_id='12345')],
+                                              statement_id='12345',
+                                              account_id='67890')],
         'Version': '2012-10-17'
     }
     lambda_stub = stubbed_session.stub('lambda')
@@ -3495,7 +3506,7 @@ def test_can_remove_s3_permission(stubbed_session):
     stubbed_session.activate_stubs()
     client = TypedAWSClient(stubbed_session)
     client.remove_permission_for_s3_event(
-        'mybucket', 'name')
+        'mybucket', 'name', '67890')
     stubbed_session.verify_stubs()
 
 

tests/unit/deploy/test_planner.py

@@ -819,6 +819,7 @@ def test_can_plan_s3_event(self):
                 params={
                     'bucket': 'mybucket',
                     'function_arn': Variable('function_name_lambda_arn'),
+                    'account_id': Variable('account_id'),
                 },
             )
         )