PLA-26855 - Updated azure remediation jobs to wait for the poller result

remediation_worker/jobs/aws_ec2_administration_ports_ingress_allowed/README.md

@@ -14,7 +14,7 @@ Network ACL should restrict administration ports (3389 and 22) from public acces
 
 ### Prerequisites
 
-The provided AWS credential must have access to `ec2:DeleteNetworkAcl`, `ec2:DescribeNetworkAcls` and `ec2:ReplaceNetworkAclEntry`.
+The provided AWS credential must have access to `ec2:CreateNetworkAclEntry`, `ec2:DeleteNetworkAclEntry`, `DescribeNetworkAcls` and `ec2:ReplaceNetworkAclEntry`.
 
 You may find the latest example policy file [here](minimum_policy.json)
 

remediation_worker/jobs/aws_ec2_administration_ports_ingress_allowed/minimum_policy.json

@@ -5,8 +5,10 @@
             "Sid": "RemoveAdministrationPortsIngress",
             "Effect": "Allow",
             "Action": [
+                "ec2:CreateNetworkAclEntry",
                 "ec2:DescribeNetworkAcls",
-                "ec2:DeleteNetworkAcl"
+                "ec2:DeleteNetworkAclEntry",
+                "ec2:ReplaceNetworkAclEntry"
             ],
             "Resource": "*"
         }

remediation_worker/jobs/azure_mysql_enforce_ssl_connection_enable/azure_mysql_enforce_ssl_connection_enable.py

@@ -77,11 +77,17 @@ def remediate(self, client, resource_group_name, mysql_server_name):
             logging.info(f"      resource_group_name={resource_group_name}")
             logging.info(f"      server_name={mysql_server_name}")
 
-            client.servers.begin_update(
+            poller = client.servers.begin_update(
                 resource_group_name=resource_group_name,
                 server_name=mysql_server_name,
                 parameters=ServerUpdateParameters(ssl_enforcement="Enabled"),
             )
+            while not poller.done():
+                time.sleep(5)
+                status = poller.status()
+                logging.info(f"The remediation job status: {status}")
+            poller.result()
+
         except Exception as e:
             logging.error(f"{str(e)}")
             raise

remediation_worker/jobs/azure_postgresql_allow_access_to_azure_service_disabled/azure_postgresql_allow_access_to_azure_service_disabled.py

@@ -70,11 +70,17 @@ def remediate(self, client, resource_group_name, postgre_server_name):
             logging.info(f"      server_name={postgre_server_name}")
             logging.info(f"      firewall_rule_name=AllowAllWindowsAzureIps")
 
-            client.firewall_rules.begin_delete(
+            poller = client.firewall_rules.begin_delete(
                 resource_group_name=resource_group_name,
                 server_name=postgre_server_name,
                 firewall_rule_name="AllowAllWindowsAzureIps",
             )
+            while not poller.done():
+                time.sleep(5)
+                status = poller.status()
+                logging.info(f"The remediation job status: {status}")
+            poller.result()
+
         except Exception as e:
             logging.error(f"{str(e)}")
             raise

remediation_worker/jobs/azure_postgresql_enforce_ssl_connection_enable/azure_postgresql_enforce_ssl_connection_enable.py

@@ -70,11 +70,17 @@ def remediate(self, client, resource_group_name, postgre_server_name):
             logging.info(f"      resource_group_name={resource_group_name}")
             logging.info(f"      server_name={postgre_server_name}")
 
-            client.servers.begin_update(
+            poller = client.servers.begin_update(
                 resource_group_name=resource_group_name,
                 server_name=postgre_server_name,
                 parameters=ServerUpdateParameters(ssl_enforcement="Enabled"),
             )
+            while not poller.done():
+                time.sleep(5)
+                status = poller.status()
+                logging.info(f"The remediation job status: {status}")
+            poller.result()
+
         except Exception as e:
             logging.error(f"{str(e)}")
             raise

remediation_worker/jobs/azure_sql_auditing_on_server/azure_sql_auditing_on_server.py

@@ -571,14 +571,19 @@ def create_server_blob_auditing_policy(
         logging.info(f"      resource_group_name={resource_group_name}")
         logging.info(f"      server_name={sql_server_name}")
 
-        client.server_blob_auditing_policies.create_or_update(
+        poller = client.server_blob_auditing_policies.create_or_update(
             resource_group_name=resource_group_name,
             server_name=sql_server_name,
             parameters=ServerBlobAuditingPolicy(
                 state=BlobAuditingPolicyState.enabled,
                 storage_endpoint=f"https://{stg_account_name}.blob.core.windows.net/",
             ),
         )
+        while not poller.done():
+            time.sleep(5)
+            status = poller.status()
+            logging.info(f"The remediation job status: {status}")
+        poller.result()
 
     def ensure_identity_assigned(
         self, client, resource_group_name, sql_server_name, region

remediation_worker/jobs/azure_sql_tde_protector_encrypted_cmk/README.md

@@ -27,6 +27,7 @@ The provided Azure service principal must have the following permissions:
 `Microsoft.KeyVault/vaults/accessPolicies/write`,
 `Microsoft.Sql/servers/read`,
 `Microsoft.Sql/servers/write`,
+`Microsoft.Sql/servers/encryptionProtector/read`,
 `Microsoft.Sql/servers/encryptionProtector/write`,
 `Microsoft.Sql/servers/keys/write`
 `Microsoft.Sql/servers/keys/read`.

remediation_worker/jobs/azure_sql_tde_protector_encrypted_cmk/azure_sql_tde_protector_encrypted_cmk.py

@@ -688,14 +688,20 @@ def remediate(
             ).result()
 
             # Update the SQL TDE protector to encrypt using cmk
-            client.encryption_protectors.begin_create_or_update(
+            poller = client.encryption_protectors.begin_create_or_update(
                 resource_group_name=resource_group_name,
                 server_name=sql_server_name,
                 encryption_protector_name="current",
                 parameters=EncryptionProtector(
                     server_key_name=server_key_name, server_key_type="AzureKeyVault"
                 ),
             )
+            while not poller.done():
+                time.sleep(5)
+                status = poller.status()
+                logging.info(f"The remediation job status: {status}")
+            poller.result()
+
         except Exception as e:
             logging.error(f"{str(e)}")
             raise

remediation_worker/jobs/azure_sql_tde_protector_encrypted_cmk/minimum_permissions.json

@@ -19,6 +19,7 @@
             "Microsoft.KeyVault/vaults/accessPolicies/write",
             "Microsoft.Sql/servers/read",
             "Microsoft.Sql/servers/write",
+            "Microsoft.Sql/servers/encryptionProtector/read",
             "Microsoft.Sql/servers/encryptionProtector/write",
             "Microsoft.Sql/servers/keys/write",
             "Microsoft.Sql/servers/keys/read"

remediation_worker/jobs/azure_sql_threat_detection_on_server/azure_sql_threat_detection_on_server.py

@@ -70,13 +70,19 @@ def remediate(self, client, resource_group_name, sql_server_name):
             logging.info(f"      resource_group_name={resource_group_name}")
             logging.info(f"      server_name={sql_server_name}")
 
-            client.server_security_alert_policies.create_or_update(
+            poller = client.server_security_alert_policies.create_or_update(
                 resource_group_name=resource_group_name,
                 server_name=sql_server_name,
                 parameters=ServerSecurityAlertPolicy(
                     state=SecurityAlertPolicyState.enabled
                 ),
             )
+            while not poller.done():
+                time.sleep(5)
+                status = poller.status()
+                logging.info(f"The remediation job status: {status}")
+            poller.result()
+
         except Exception as e:
             logging.error(f"{str(e)}")
             raise

remediation_worker/jobs/azure_sql_threat_detection_types_all_server/azure_sql_threat_detection_types_all_server.py

@@ -73,13 +73,19 @@ def remediate(self, client, resource_group_name, sql_server_name):
             logging.info(f"      resource_group_name={resource_group_name}")
             logging.info(f"      server_name={sql_server_name}")
 
-            client.server_security_alert_policies.create_or_update(
+            poller = client.server_security_alert_policies.create_or_update(
                 resource_group_name=resource_group_name,
                 server_name=sql_server_name,
                 parameters=ServerSecurityAlertPolicy(
                     state=SecurityAlertPolicyState.enabled, disabled_alerts=[]
                 ),
             )
+            while not poller.done():
+                time.sleep(5)
+                status = poller.status()
+                logging.info(f"The remediation job status: {status}")
+            poller.result()
+
         except Exception as e:
             logging.error(f"{str(e)}")
             raise

test/unit/test_azure_sql_auditing_on_server.py

@@ -77,6 +77,7 @@ def test_remediate_without_stg_without_keyvault(self):
         action.ensure_identity_assigned = Mock()
         action.create_role_assignment = Mock()
         action.create_server_blob_auditing_policy = Mock()
+        action.check_role_assignment = Mock()
 
         identity = Identity(
             principal_id="139bcf82-e14e-4773-bcf4-1da136674792",
@@ -106,6 +107,7 @@ def test_remediate_without_stg_without_keyvault(self):
             location="eastus",
             identity=identity,
         )
+        action.check_role_assignment.return_value = True
         action.check_stg_account.return_value = None
         action.check_key_vault.return_value = None
         assert (
@@ -134,7 +136,6 @@ def test_remediate_without_stg_without_keyvault(self):
         assert action.create_storage_account.call_count == 1
         assert action.create_key_vault.call_count == 1
         assert action.create_diagnostic_setting.call_count == 1
-        assert action.create_role_assignment.call_count == 1
         assert action.create_server_blob_auditing_policy.call_count == 1
 
     def test_remediate_without_stg_with_keyvault(self):
@@ -159,6 +160,7 @@ def test_remediate_without_stg_with_keyvault(self):
         action.ensure_identity_assigned = Mock()
         action.create_role_assignment = Mock()
         action.create_server_blob_auditing_policy = Mock()
+        action.check_role_assignment = Mock()
 
         identity = Identity(
             principal_id="139bcf82-e14e-4773-bcf4-1da136674792",
@@ -189,6 +191,7 @@ def test_remediate_without_stg_with_keyvault(self):
                 vault_uri="https://stg-keyvault-rem.vault.azure.net",
             ),
         )
+        action.check_role_assignment.return_value = False
         assert (
             action.remediate(
                 client_id,