From 114b567ee566366c6ee53388f3c53482ee05189c Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Tue, 13 Feb 2024 10:41:21 +0100
Subject: [PATCH] Added drop table and all permissions to pivot role as safety
 measure for actions in revoke share

---
 .../share_managers/lf_share_manager.py        | 15 +++++++++++
 .../lakeformation_process_share.py            |  2 ++
 .../datasets/tasks/test_lf_share_manager.py   | 25 +++++++++++++++++++
 3 files changed, 42 insertions(+)

diff --git a/backend/dataall/modules/dataset_sharing/services/share_managers/lf_share_manager.py b/backend/dataall/modules/dataset_sharing/services/share_managers/lf_share_manager.py
index 94e629238..b3f2eab95 100644
--- a/backend/dataall/modules/dataset_sharing/services/share_managers/lf_share_manager.py
+++ b/backend/dataall/modules/dataset_sharing/services/share_managers/lf_share_manager.py
@@ -202,6 +202,21 @@ def grant_pivot_role_all_database_permissions_to_shared_database(self) -> True:
         )
         return True
 
+    def grant_pivot_role_drop_permissions_to_resource_link_table(self, table: DatasetTable) -> True:
+        """
+        Grants 'DROP' Lake Formation permissions to pivot role to the resource link table in target account
+        :param table: DatasetTable
+        :return: True if it is successful
+        """
+        self.lf_client_in_target.grant_permissions_to_table(
+            principals=[SessionHelper.get_delegation_role_arn(self.target_environment.AwsAccountId)],
+            database_name=self.shared_db_name,
+            table_name=table.GlueTableName,
+            catalog_id=self.target_environment.AwsAccountId,
+            permissions=['DROP']
+        )
+        return True
+
     def grant_principals_database_permissions_to_shared_database(self) -> True:
         """
         Grants 'DESCRIBE' Lake Formation permissions to share principals to the shared database in target account
diff --git a/backend/dataall/modules/dataset_sharing/services/share_processors/lakeformation_process_share.py b/backend/dataall/modules/dataset_sharing/services/share_processors/lakeformation_process_share.py
index 1406e6eaf..f60a57e02 100644
--- a/backend/dataall/modules/dataset_sharing/services/share_processors/lakeformation_process_share.py
+++ b/backend/dataall/modules/dataset_sharing/services/share_processors/lakeformation_process_share.py
@@ -155,6 +155,7 @@ def process_revoked_shares(self) -> bool:
             '##### Starting Revoking tables #######'
         )
         success = True
+        self.grant_pivot_role_all_database_permissions_to_shared_database()
         for table in self.revoked_tables:
             share_item = ShareObjectRepository.find_sharable_item(
                 self.session, self.share.shareUri, table.tableUri
@@ -182,6 +183,7 @@ def process_revoked_shares(self) -> bool:
 
                     if (self.is_new_share and not other_table_shares_in_env) or not self.is_new_share:
                         warn('self.is_new_share will be deprecated in v2.6.0', DeprecationWarning, stacklevel=2)
+                        self.grant_pivot_role_drop_permissions_to_resource_link_table(table)
                         self.delete_resource_link_table_in_shared_database(table)
 
                 if not other_table_shares_in_env:
diff --git a/tests/modules/datasets/tasks/test_lf_share_manager.py b/tests/modules/datasets/tasks/test_lf_share_manager.py
index 54a39ca85..cb7c48de6 100644
--- a/tests/modules/datasets/tasks/test_lf_share_manager.py
+++ b/tests/modules/datasets/tasks/test_lf_share_manager.py
@@ -445,6 +445,31 @@ def test_grant_principals_permissions_to_resource_link_table(
             permissions=['DESCRIBE']
         )
 
+
+def test_grant_pivot_role_drop_permissions_to_resource_link_table(
+        processor_with_mocks,
+        table1: DatasetTable,
+        target_environment: Environment,
+        mocker
+):
+    processor, lf_client, glue_client = processor_with_mocks
+    mocker.patch(
+        "dataall.base.aws.sts.SessionHelper.get_delegation_role_arn",
+        return_value="arn:role",
+    )
+    # When
+    processor.grant_pivot_role_drop_permissions_to_resource_link_table(table1)
+    # Then
+    lf_client.grant_permissions_to_table.assert_called_once()
+    lf_client.grant_permissions_to_table.assert_called_with(
+            principals=["arn:role"],
+            database_name=processor.shared_db_name,
+            table_name=table1.GlueTableName,
+            catalog_id=target_environment.AwsAccountId,
+            permissions=['DROP']
+        )
+
+
 def test_grant_principals_permissions_to_table_in_target(
         processor_with_mocks,
         table1: DatasetTable,