SSM: add support for maintenance window tasks (#6430)

IMPLEMENTATION_COVERAGE.md

@@ -6784,7 +6784,7 @@
 - [ ] deregister_managed_instance
 - [ ] deregister_patch_baseline_for_patch_group
 - [X] deregister_target_from_maintenance_window
-- [ ] deregister_task_from_maintenance_window
+- [X] deregister_task_from_maintenance_window
 - [ ] describe_activations
 - [ ] describe_association
 - [ ] describe_association_execution_targets
@@ -6807,7 +6807,7 @@
 - [ ] describe_maintenance_window_executions
 - [ ] describe_maintenance_window_schedule
 - [X] describe_maintenance_window_targets
-- [ ] describe_maintenance_window_tasks
+- [X] describe_maintenance_window_tasks
 - [X] describe_maintenance_windows
 - [ ] describe_maintenance_windows_for_target
 - [ ] describe_ops_items
@@ -6868,7 +6868,7 @@
 - [ ] register_default_patch_baseline
 - [ ] register_patch_baseline_for_patch_group
 - [X] register_target_with_maintenance_window
-- [ ] register_task_with_maintenance_window
+- [X] register_task_with_maintenance_window
 - [X] remove_tags_from_resource
 - [ ] reset_service_setting
 - [ ] resume_session

moto/ssm/models.py

@@ -975,6 +975,65 @@ class FakeMaintenanceWindowTarget:
         return str(random.uuid4())
 
 
+class FakeMaintenanceWindowTask:
+    def __init__(
+        self,
+        window_id: str,
+        targets: Optional[List[Dict[str, Any]]],
+        task_arn: str,
+        service_role_arn: Optional[str],
+        task_type: str,
+        task_parameters: Optional[Dict[str, Any]],
+        task_invocation_parameters: Optional[Dict[str, Any]],
+        priority: Optional[int],
+        max_concurrency: Optional[str],
+        max_errors: Optional[str],
+        logging_info: Optional[Dict[str, Any]],
+        name: Optional[str],
+        description: Optional[str],
+        cutoff_behavior: Optional[str],
+        alarm_configurations: Optional[Dict[str, Any]],
+    ):
+        self.window_task_id = FakeMaintenanceWindowTask.generate_id()
+        self.window_id = window_id
+        self.task_type = task_type
+        self.task_arn = task_arn
+        self.service_role_arn = service_role_arn
+        self.task_parameters = task_parameters
+        self.priority = priority
+        self.max_concurrency = max_concurrency
+        self.max_errors = max_errors
+        self.logging_info = logging_info
+        self.name = name
+        self.description = description
+        self.targets = targets
+        self.task_invocation_parameters = task_invocation_parameters
+        self.cutoff_behavior = cutoff_behavior
+        self.alarm_configurations = alarm_configurations
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "WindowId": self.window_id,
+            "WindowTaskId": self.window_task_id,
+            "TaskType": self.task_type,
+            "TaskArn": self.task_arn,
+            "ServiceRoleArn": self.service_role_arn,
+            "TaskParameters": self.task_parameters,
+            "Priority": self.priority,
+            "MaxConcurrency": self.max_concurrency,
+            "MaxErrors": self.max_errors,
+            "LoggingInfo": self.logging_info,
+            "Name": self.name,
+            "Description": self.description,
+            "Targets": self.targets,
+            "TaskInvocationParameters": self.task_invocation_parameters,
+        }
+
+    @staticmethod
+    def generate_id() -> str:
+        return str(random.uuid4())
+
+
 def _maintenance_window_target_filter_match(
     filters: Optional[List[Dict[str, Any]]], target: FakeMaintenanceWindowTarget
 ) -> bool:
@@ -983,6 +1042,14 @@ def _maintenance_window_target_filter_match(
     return False
 
 
+def _maintenance_window_task_filter_match(
+    filters: Optional[List[Dict[str, Any]]], task: FakeMaintenanceWindowTask
+) -> bool:
+    if not filters and task:
+        return True
+    return False
+
+
 class FakeMaintenanceWindow:
     def __init__(
         self,
@@ -1007,6 +1074,7 @@ class FakeMaintenanceWindow:
         self.start_date = start_date
         self.end_date = end_date
         self.targets: Dict[str, FakeMaintenanceWindowTarget] = {}
+        self.tasks: Dict[str, FakeMaintenanceWindowTask] = {}
 
     def to_json(self) -> Dict[str, Any]:
         return {
@@ -2351,5 +2419,62 @@ class SimpleSystemManagerBackend(BaseBackend):
         ]
         return targets
 
+    def register_task_with_maintenance_window(
+        self,
+        window_id: str,
+        targets: Optional[List[Dict[str, Any]]],
+        task_arn: str,
+        service_role_arn: Optional[str],
+        task_type: str,
+        task_parameters: Optional[Dict[str, Any]],
+        task_invocation_parameters: Optional[Dict[str, Any]],
+        priority: Optional[int],
+        max_concurrency: Optional[str],
+        max_errors: Optional[str],
+        logging_info: Optional[Dict[str, Any]],
+        name: Optional[str],
+        description: Optional[str],
+        cutoff_behavior: Optional[str],
+        alarm_configurations: Optional[Dict[str, Any]],
+    ) -> str:
+
+        window = self.get_maintenance_window(window_id)
+        task = FakeMaintenanceWindowTask(
+            window_id,
+            targets,
+            task_arn,
+            service_role_arn,
+            task_type,
+            task_parameters,
+            task_invocation_parameters,
+            priority,
+            max_concurrency,
+            max_errors,
+            logging_info,
+            name,
+            description,
+            cutoff_behavior,
+            alarm_configurations,
+        )
+        window.tasks[task.window_task_id] = task
+        return task.window_task_id
+
+    def describe_maintenance_window_tasks(
+        self, window_id: str, filters: List[Dict[str, Any]]
+    ) -> List[FakeMaintenanceWindowTask]:
+        window = self.get_maintenance_window(window_id)
+        tasks = [
+            task
+            for task in window.tasks.values()
+            if _maintenance_window_task_filter_match(filters, task)
+        ]
+        return tasks
+
+    def deregister_task_from_maintenance_window(
+        self, window_id: str, window_task_id: str
+    ) -> None:
+        window = self.get_maintenance_window(window_id)
+        del window.tasks[window_task_id]
+
 
 ssm_backends = BackendDict(SimpleSystemManagerBackend, "ssm")

moto/ssm/responses.py

@@ -513,3 +513,42 @@ class SimpleSystemManagerResponse(BaseResponse):
         baseline_id = self._get_param("BaselineId")
         self.ssm_backend.delete_patch_baseline(baseline_id)
         return "{}"
+
+    def register_task_with_maintenance_window(self) -> str:
+        window_task_id = self.ssm_backend.register_task_with_maintenance_window(
+            window_id=self._get_param("WindowId"),
+            targets=self._get_param("Targets"),
+            task_arn=self._get_param("TaskArn"),
+            service_role_arn=self._get_param("ServiceRoleArn"),
+            task_type=self._get_param("TaskType"),
+            task_parameters=self._get_param("TaskParameters"),
+            task_invocation_parameters=self._get_param("TaskInvocationParameters"),
+            priority=self._get_param("Priority"),
+            max_concurrency=self._get_param("MaxConcurrency"),
+            max_errors=self._get_param("MaxErrors"),
+            logging_info=self._get_param("LoggingInfo"),
+            name=self._get_param("Name"),
+            description=self._get_param("Description"),
+            cutoff_behavior=self._get_param("CutoffBehavior"),
+            alarm_configurations=self._get_param("AlarmConfigurations"),
+        )
+        return json.dumps({"WindowTaskId": window_task_id})
+
+    def describe_maintenance_window_tasks(self) -> str:
+        window_id = self._get_param("WindowId")
+        filters = self._get_param("Filters", [])
+        tasks = [
+            task.to_json()
+            for task in self.ssm_backend.describe_maintenance_window_tasks(
+                window_id, filters
+            )
+        ]
+        return json.dumps({"Tasks": tasks})
+
+    def deregister_task_from_maintenance_window(self) -> str:
+        window_id = self._get_param("WindowId")
+        window_task_id = self._get_param("WindowTaskId")
+        self.ssm_backend.deregister_task_from_maintenance_window(
+            window_id, window_task_id
+        )
+        return "{}"

tests/test_ssm/test_ssm_maintenance_windows.py

@@ -274,3 +274,111 @@ def test_deregister_target_from_maintenance_window():
         WindowId=window_id,
     )
     resp.should.have.key("Targets").should.have.length_of(0)
+
+
+@mock_ssm
+def test_describe_maintenance_window_with_no_task_or_targets():
+    ssm = boto3.client("ssm", region_name="us-east-1")
+
+    resp = ssm.create_maintenance_window(
+        Name="simple-window",
+        Schedule="cron(15 12 * * ? *)",
+        Duration=2,
+        Cutoff=1,
+        AllowUnassociatedTargets=False,
+    )
+    window_id = resp["WindowId"]
+
+    resp = ssm.describe_maintenance_window_tasks(
+        WindowId=window_id,
+    )
+    resp.should.have.key("Tasks").should.have.length_of(0)
+
+    resp = ssm.describe_maintenance_window_targets(
+        WindowId=window_id,
+    )
+    resp.should.have.key("Targets").should.have.length_of(0)
+
+
+@mock_ssm
+def test_register_maintenance_window_task():
+    ssm = boto3.client("ssm", region_name="us-east-1")
+
+    resp = ssm.create_maintenance_window(
+        Name="simple-window",
+        Schedule="cron(15 12 * * ? *)",
+        Duration=2,
+        Cutoff=1,
+        AllowUnassociatedTargets=False,
+    )
+    window_id = resp["WindowId"]
+
+    resp = ssm.register_target_with_maintenance_window(
+        WindowId=window_id,
+        ResourceType="INSTANCE",
+        Targets=[{"Key": "tag:Name", "Values": ["my-instance"]}],
+    )
+    window_target_id = resp["WindowTargetId"]
+
+    resp = ssm.register_task_with_maintenance_window(
+        WindowId=window_id,
+        Targets=[{"Key": "WindowTargetIds", "Values": [window_target_id]}],
+        TaskArn="AWS-RunShellScript",
+        TaskType="RUN_COMMAND",
+        MaxConcurrency="1",
+        MaxErrors="1",
+    )
+
+    resp.should.have.key("WindowTaskId")
+    _id = resp["WindowTaskId"]
+
+    resp = ssm.describe_maintenance_window_tasks(
+        WindowId=window_id,
+    )
+    resp.should.have.key("Tasks").should.have.length_of(1)
+    resp["Tasks"][0].should.have.key("WindowTaskId").equal(_id)
+    resp["Tasks"][0].should.have.key("WindowId").equal(window_id)
+    resp["Tasks"][0].should.have.key("TaskArn").equal("AWS-RunShellScript")
+    resp["Tasks"][0].should.have.key("MaxConcurrency").equal("1")
+    resp["Tasks"][0].should.have.key("MaxErrors").equal("1")
+
+
+@mock_ssm
+def test_deregister_maintenance_window_task():
+    ssm = boto3.client("ssm", region_name="us-east-1")
+
+    resp = ssm.create_maintenance_window(
+        Name="simple-window",
+        Schedule="cron(15 12 * * ? *)",
+        Duration=2,
+        Cutoff=1,
+        AllowUnassociatedTargets=False,
+    )
+    window_id = resp["WindowId"]
+
+    resp = ssm.register_target_with_maintenance_window(
+        WindowId=window_id,
+        ResourceType="INSTANCE",
+        Targets=[{"Key": "tag:Name", "Values": ["my-instance"]}],
+    )
+    window_target_id = resp["WindowTargetId"]
+
+    resp = ssm.register_task_with_maintenance_window(
+        WindowId=window_id,
+        Targets=[{"Key": "WindowTargetIds", "Values": [window_target_id]}],
+        TaskArn="AWS-RunShellScript",
+        TaskType="RUN_COMMAND",
+        MaxConcurrency="1",
+        MaxErrors="1",
+    )
+    window_task_id = resp["WindowTaskId"]
+
+    ssm.deregister_task_from_maintenance_window(
+        WindowId=window_id,
+        WindowTaskId=window_task_id,
+    )
+
+    resp = ssm.describe_maintenance_window_tasks(
+        WindowId=window_id,
+    )
+    resp.should.have.key("Tasks").should.have.length_of(0)