Repository: cloud-custodian/cloud-custodian

Issue:
Airflow - Add update and delete actions
### Describe the feature

I'd like to be able to use a custodian action to enforce the use of `PRIVATE_ONLY` for the `WebserverAccessMode`. Ideally, something like an `update-environment` action might be able to be used as a corrective action but if that's not feasible, a `delete-environment` action could nuke the resource. 

### Extra information or context

I will be working on this issue


Discussion:
You could potentially look at adding an `update-environment` action that would include `WebserverAccessMode` settings. A delete action here would also make sense but for the sake of corrective action, update seems more appropriate

Resolving patch:
diff --git a/c7n/resources/airflow.py b/c7n/resources/airflow.py
--- a/c7n/resources/airflow.py
+++ b/c7n/resources/airflow.py
@@ -4,6 +4,8 @@
 from c7n.query import QueryResourceManager, TypeInfo
 from c7n.filters.kms import KmsRelatedFilter
 from c7n.tags import RemoveTag, Tag, TagDelayedAction, TagActionFilter
+from c7n.actions import Action
+from c7n.utils import local_session, type_schema
 
 
 @resources.register('airflow')
@@ -113,3 +115,77 @@ def process_resource_set(self, client, airflow, tags):
 
 ApacheAirflow.filter_registry.register('marked-for-op', TagActionFilter)
 ApacheAirflow.action_registry.register('mark-for-op', TagDelayedAction)
+
+@ApacheAirflow.action_registry.register('update-environment')
+class UpdateApacheAirflowEnvironment(Action):
+    """
+    Action to update an Airflow environment to
+    set the WebserverAccessMode to PRIVATE_ONLY or PUBLIC_ONLY.
+
+    :example:
+
+    .. code-block:: yaml
+
+        policies:
+          - name: set-webserver-access-mode
+            resource: airflow
+            actions:
+              - type: update-environment
+                access_mode: PRIVATE_ONLY
+    """
+
+    permissions = ('airflow:UpdateEnvironment',)
+
+    schema = type_schema(
+        'update-environment',
+        access_mode={'type': 'string', 'enum': ['PRIVATE_ONLY', 'PUBLIC_ONLY']},
+        required=['access_mode']
+    )
+    valid_origin_states = ('AVAILABLE', 'UPDATE_FAILED')
+
+    def process(self, resources):
+        resources = self.filter_resources(resources, 'Status', self.valid_origin_states)
+        client = local_session(self.manager.session_factory).client('mwaa')
+        access_mode = self.data.get('access_mode')
+        for r in resources:
+          self.process_environment(r, client, access_mode)
+
+    def process_environment(self, r, client, access_mode):
+        current_access_mode = r.get('WebserverAccessMode')
+        if current_access_mode == access_mode:
+          return
+        client.update_environment(
+          Name=r['Name'],
+          WebserverAccessMode=access_mode
+        )
+
+@ApacheAirflow.action_registry.register('delete-environment')
+class DeleteApacheAirflowEnvironment(Action):
+    """
+    Action to delete a Managed Workflow for Apache Airflow environment
+
+    :example:
+
+    .. code-block:: yaml
+
+        policies:
+          - name: delete-airflow-environment
+            resource: airflow
+            actions:
+              - type: delete-environment
+    """
+
+    permissions = ('airflow:DeleteEnvironment',)
+
+    schema = type_schema('delete-environment')
+    valid_origin_states = ('AVAILABLE', 'CREATE_FAILED', 'DELETE_FAILED',)
+
+    def process(self, resources):
+        resources = self.filter_resources(resources, 'Status', self.valid_origin_states)
+        client = local_session(self.manager.session_factory).client('mwaa')
+        for r in resources:
+          self.manager.retry(
+            client.delete_environment,
+            Name=r["Name"],
+            ignore_err_codes=("ResourceNotFoundException",)
+          )


Test patch:
diff --git a/tests/data/placebo/test_airflow_delete_environment/airflow.DeleteEnvironment_1.json b/tests/data/placebo/test_airflow_delete_environment/airflow.DeleteEnvironment_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_airflow_delete_environment/airflow.DeleteEnvironment_1.json
@@ -0,0 +1,6 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {}
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_airflow_delete_environment/airflow.GetEnvironment_1.json b/tests/data/placebo/test_airflow_delete_environment/airflow.GetEnvironment_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_airflow_delete_environment/airflow.GetEnvironment_1.json
@@ -0,0 +1,83 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {},
+        "Environment": {
+            "AirflowVersion": "2.6.3",
+            "Arn": "arn:aws:airflow:us-east-1:644160558196:environment/TestStack-MwaaEnvironment",
+            "CreatedAt": {
+                "__class__": "datetime",
+                "year": 2023,
+                "month": 9,
+                "day": 13,
+                "hour": 6,
+                "minute": 43,
+                "second": 40,
+                "microsecond": 0
+            },
+            "DagS3Path": "dags",
+            "EnvironmentClass": "mw1.small",
+            "ExecutionRoleArn": "arn:aws:iam::644160558196:role/service-role/TestStack-MwaaExecutionRole-J13JXICOSM01",
+            "LastUpdate": {
+                "CreatedAt": {
+                    "__class__": "datetime",
+                    "year": 2023,
+                    "month": 9,
+                    "day": 13,
+                    "hour": 17,
+                    "minute": 29,
+                    "second": 9,
+                    "microsecond": 0
+                },
+                "Status": "SUCCESS"
+            },
+            "LoggingConfiguration": {
+                "DagProcessingLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-DAGProcessing",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "SchedulerLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Scheduler",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "TaskLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Task",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "WebserverLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-WebServer",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "WorkerLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Worker",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                }
+            },
+            "MaxWorkers": 2,
+            "MinWorkers": 1,
+            "Name": "TestStack-MwaaEnvironment",
+            "NetworkConfiguration": {
+                "SecurityGroupIds": [
+                    "sg-014d3b1a314ccede1"
+                ],
+                "SubnetIds": [
+                    "subnet-0e4d3bc0b5764db4f",
+                    "subnet-040942ec3da9ce524"
+                ]
+            },
+            "Schedulers": 2,
+            "ServiceRoleArn": "arn:aws:iam::644160558196:role/aws-service-role/airflow.amazonaws.com/AWSServiceRoleForAmazonMWAA",
+            "SourceBucketArn": "arn:aws:s3:::teststack-environmentbucket-f4ymk3qn4gj4",
+            "Status": "AVAILABLE",
+            "Tags": {},
+            "WebserverAccessMode": "PRIVATE_ONLY",
+            "WebserverUrl": "915dda6f-4589-4d17-8228-5fb093cecf27-vpce.c61.us-east-1.airflow.amazonaws.com",
+            "WeeklyMaintenanceWindowStart": "WED:18:30"
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_airflow_delete_environment/airflow.GetEnvironment_2.json b/tests/data/placebo/test_airflow_delete_environment/airflow.GetEnvironment_2.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_airflow_delete_environment/airflow.GetEnvironment_2.json
@@ -0,0 +1,83 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {},
+        "Environment": {
+            "AirflowVersion": "2.6.3",
+            "Arn": "arn:aws:airflow:us-east-1:644160558196:environment/TestStack-MwaaEnvironment",
+            "CreatedAt": {
+                "__class__": "datetime",
+                "year": 2023,
+                "month": 9,
+                "day": 13,
+                "hour": 6,
+                "minute": 43,
+                "second": 40,
+                "microsecond": 0
+            },
+            "DagS3Path": "dags",
+            "EnvironmentClass": "mw1.small",
+            "ExecutionRoleArn": "arn:aws:iam::644160558196:role/service-role/TestStack-MwaaExecutionRole-J13JXICOSM01",
+            "LastUpdate": {
+                "CreatedAt": {
+                    "__class__": "datetime",
+                    "year": 2023,
+                    "month": 9,
+                    "day": 13,
+                    "hour": 17,
+                    "minute": 29,
+                    "second": 9,
+                    "microsecond": 0
+                },
+                "Status": "SUCCESS"
+            },
+            "LoggingConfiguration": {
+                "DagProcessingLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-DAGProcessing",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "SchedulerLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Scheduler",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "TaskLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Task",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "WebserverLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-WebServer",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "WorkerLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Worker",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                }
+            },
+            "MaxWorkers": 2,
+            "MinWorkers": 1,
+            "Name": "TestStack-MwaaEnvironment",
+            "NetworkConfiguration": {
+                "SecurityGroupIds": [
+                    "sg-014d3b1a314ccede1"
+                ],
+                "SubnetIds": [
+                    "subnet-0e4d3bc0b5764db4f",
+                    "subnet-040942ec3da9ce524"
+                ]
+            },
+            "Schedulers": 2,
+            "ServiceRoleArn": "arn:aws:iam::644160558196:role/aws-service-role/airflow.amazonaws.com/AWSServiceRoleForAmazonMWAA",
+            "SourceBucketArn": "arn:aws:s3:::teststack-environmentbucket-f4ymk3qn4gj4",
+            "Status": "DELETING",
+            "Tags": {},
+            "WebserverAccessMode": "PRIVATE_ONLY",
+            "WebserverUrl": "915dda6f-4589-4d17-8228-5fb093cecf27-vpce.c61.us-east-1.airflow.amazonaws.com",
+            "WeeklyMaintenanceWindowStart": "WED:18:30"
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_airflow_delete_environment/airflow.ListEnvironments_1.json b/tests/data/placebo/test_airflow_delete_environment/airflow.ListEnvironments_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_airflow_delete_environment/airflow.ListEnvironments_1.json
@@ -0,0 +1,9 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {},
+        "Environments": [
+            "TestStack-MwaaEnvironment"
+        ]
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_airflow_update_environment/airflow.GetEnvironment_1.json b/tests/data/placebo/test_airflow_update_environment/airflow.GetEnvironment_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_airflow_update_environment/airflow.GetEnvironment_1.json
@@ -0,0 +1,83 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {},
+        "Environment": {
+            "AirflowVersion": "2.6.3",
+            "Arn": "arn:aws:airflow:us-east-1:644160558196:environment/TestStack-MwaaEnvironment",
+            "CreatedAt": {
+                "__class__": "datetime",
+                "year": 2023,
+                "month": 9,
+                "day": 14,
+                "hour": 8,
+                "minute": 28,
+                "second": 10,
+                "microsecond": 0
+            },
+            "DagS3Path": "dags",
+            "EnvironmentClass": "mw1.small",
+            "ExecutionRoleArn": "arn:aws:iam::644160558196:role/service-role/TestStack-MwaaExecutionRole-SBLUFN7YFHRU",
+            "LastUpdate": {
+                "CreatedAt": {
+                    "__class__": "datetime",
+                    "year": 2023,
+                    "month": 9,
+                    "day": 14,
+                    "hour": 8,
+                    "minute": 28,
+                    "second": 10,
+                    "microsecond": 0
+                },
+                "Status": "SUCCESS"
+            },
+            "LoggingConfiguration": {
+                "DagProcessingLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-DAGProcessing",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "SchedulerLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Scheduler",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "TaskLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Task",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "WebserverLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-WebServer",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "WorkerLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Worker",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                }
+            },
+            "MaxWorkers": 2,
+            "MinWorkers": 1,
+            "Name": "TestStack-MwaaEnvironment",
+            "NetworkConfiguration": {
+                "SecurityGroupIds": [
+                    "sg-038d1f22f343477a6"
+                ],
+                "SubnetIds": [
+                    "subnet-0a1d5260c72c1969c",
+                    "subnet-0870aa18157ea3f08"
+                ]
+            },
+            "Schedulers": 2,
+            "ServiceRoleArn": "arn:aws:iam::644160558196:role/aws-service-role/airflow.amazonaws.com/AWSServiceRoleForAmazonMWAA",
+            "SourceBucketArn": "arn:aws:s3:::teststack-environmentbucket-1uby1afpuyeec",
+            "Status": "AVAILABLE",
+            "Tags": {},
+            "WebserverAccessMode": "PUBLIC_ONLY",
+            "WebserverUrl": "836a0ccd-034c-4386-bdb2-938be4ee0454.c21.us-east-1.airflow.amazonaws.com",
+            "WeeklyMaintenanceWindowStart": "THU:18:00"
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_airflow_update_environment/airflow.GetEnvironment_2.json b/tests/data/placebo/test_airflow_update_environment/airflow.GetEnvironment_2.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_airflow_update_environment/airflow.GetEnvironment_2.json
@@ -0,0 +1,83 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {},
+        "Environment": {
+            "AirflowVersion": "2.6.3",
+            "Arn": "arn:aws:airflow:us-east-1:644160558196:environment/TestStack-MwaaEnvironment",
+            "CreatedAt": {
+                "__class__": "datetime",
+                "year": 2023,
+                "month": 9,
+                "day": 14,
+                "hour": 8,
+                "minute": 28,
+                "second": 10,
+                "microsecond": 0
+            },
+            "DagS3Path": "dags",
+            "EnvironmentClass": "mw1.small",
+            "ExecutionRoleArn": "arn:aws:iam::644160558196:role/service-role/TestStack-MwaaExecutionRole-SBLUFN7YFHRU",
+            "LastUpdate": {
+                "CreatedAt": {
+                    "__class__": "datetime",
+                    "year": 2023,
+                    "month": 9,
+                    "day": 14,
+                    "hour": 9,
+                    "minute": 5,
+                    "second": 5,
+                    "microsecond": 0
+                },
+                "Status": "PENDING"
+            },
+            "LoggingConfiguration": {
+                "DagProcessingLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-DAGProcessing",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "SchedulerLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Scheduler",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "TaskLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Task",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "WebserverLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-WebServer",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                },
+                "WorkerLogs": {
+                    "CloudWatchLogGroupArn": "arn:aws:logs:us-east-1:644160558196:log-group:airflow-TestStack-MwaaEnvironment-Worker",
+                    "Enabled": true,
+                    "LogLevel": "INFO"
+                }
+            },
+            "MaxWorkers": 2,
+            "MinWorkers": 1,
+            "Name": "TestStack-MwaaEnvironment",
+            "NetworkConfiguration": {
+                "SecurityGroupIds": [
+                    "sg-038d1f22f343477a6"
+                ],
+                "SubnetIds": [
+                    "subnet-0a1d5260c72c1969c",
+                    "subnet-0870aa18157ea3f08"
+                ]
+            },
+            "Schedulers": 2,
+            "ServiceRoleArn": "arn:aws:iam::644160558196:role/aws-service-role/airflow.amazonaws.com/AWSServiceRoleForAmazonMWAA",
+            "SourceBucketArn": "arn:aws:s3:::teststack-environmentbucket-1uby1afpuyeec",
+            "Status": "AVAILABLE",
+            "Tags": {},
+            "WebserverAccessMode": "PRIVATE_ONLY",
+            "WebserverUrl": "836a0ccd-034c-4386-bdb2-938be4ee0454.c21.us-east-1.airflow.amazonaws.com",
+            "WeeklyMaintenanceWindowStart": "THU:18:00"
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_airflow_update_environment/airflow.ListEnvironments_1.json b/tests/data/placebo/test_airflow_update_environment/airflow.ListEnvironments_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_airflow_update_environment/airflow.ListEnvironments_1.json
@@ -0,0 +1,9 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {},
+        "Environments": [
+            "TestStack-MwaaEnvironment"
+        ]
+    }
+}
\ No newline at end of file
diff --git a/tests/data/placebo/test_airflow_update_environment/airflow.UpdateEnvironment_1.json b/tests/data/placebo/test_airflow_update_environment/airflow.UpdateEnvironment_1.json
new file mode 100644
--- /dev/null
+++ b/tests/data/placebo/test_airflow_update_environment/airflow.UpdateEnvironment_1.json
@@ -0,0 +1,7 @@
+{
+    "status_code": 200,
+    "data": {
+        "ResponseMetadata": {},
+        "Arn": "arn:aws:airflow:us-east-1:644160558196:environment/TestStack-MwaaEnvironment"
+    }
+}
\ No newline at end of file
diff --git a/tests/test_airflow.py b/tests/test_airflow.py
--- a/tests/test_airflow.py
+++ b/tests/test_airflow.py
@@ -95,3 +95,62 @@ def test_airflow_environment_untag(self):
         airflow = session_factory().client('mwaa')
         call = airflow.get_environment(Name=name)
         self.assertEqual({}, call['Environment'].get('Tags'))
+
+    def test_airflow_update_environment(self):
+        session_factory = self.replay_flight_data('test_airflow_update_environment')
+        p = self.load_policy(
+            {
+                "name": "airflow-update-webserver-access-mode",
+                "resource": "airflow",
+                "filters": [
+                    {
+                        "type": "value",
+                        "key": "WebserverAccessMode",
+                        "op": "eq",
+                        "value": "PUBLIC_ONLY",
+                    }
+                ],
+                "actions": [
+                    {
+                        "type": "update-environment",
+                        "access_mode": "PRIVATE_ONLY",
+                    }
+                ]
+            },
+            session_factory=session_factory,
+        )
+        resources = p.run()
+        self.assertEqual(1, len(resources))
+        name = resources[0].get('Name')
+        airflow = session_factory().client('mwaa')
+        call = airflow.get_environment(Name=name)
+        self.assertEqual("PRIVATE_ONLY", call['Environment'].get('WebserverAccessMode'))
+
+    def test_airflow_delete_environment(self):
+        session_factory = self.replay_flight_data('test_airflow_delete_environment')
+        p = self.load_policy(
+            {
+                "name": "airflow-update-webserver-access-mode",
+                "resource": "airflow",
+                "filters": [
+                    {
+                        "type": "value",
+                        "key": "Name",
+                        "op": "eq",
+                        "value": "TestStack-MwaaEnvironment",
+                    }
+                ],
+                "actions": [
+                    {
+                        "type": "delete-environment"
+                    }
+                ]
+            },
+            session_factory=session_factory,
+        )
+        resources = p.run()
+        self.assertEqual(len(resources), 1)
+        name = resources[0].get('Name')
+        airflow = session_factory().client('mwaa')
+        call = airflow.get_environment(Name=name)
+        self.assertEqual("DELETING", call['Environment'].get('Status'))