diff --git a/src/access_control.py b/src/access_control.py index 15625aa..f71cff2 100644 --- a/src/access_control.py +++ b/src/access_control.py @@ -160,6 +160,7 @@ def execute_decision( # noqa: PLR0913 approver: entities.slack.User, requester: entities.slack.User, reason: str, + thread_ts: str | None = None, ) -> bool: logger.info("Executing decision") if not decision.grant: @@ -215,6 +216,7 @@ def execute_decision( # noqa: PLR0913 permission_set_arn=permission_set.arn, user_principal_id=sso_user_principal_id, ), + thread_ts=thread_ts, ) return True # Temporary solution for testing @@ -227,6 +229,7 @@ def execute_decision_on_group_request( # noqa: PLR0913 requester: entities.slack.User, reason: str, identity_store_id: str, + thread_ts: str | None = None, ) -> bool: logger.info("Executing decision") if not decision.grant: @@ -284,5 +287,6 @@ def execute_decision_on_group_request( # noqa: PLR0913 user_principal_id=sso_user_principal_id, membership_id=membership_id, ), + thread_ts=thread_ts, ) return # type: ignore # noqa: PGH003 diff --git a/src/events.py b/src/events.py index 0daf5fd..55b1bec 100644 --- a/src/events.py +++ b/src/events.py @@ -14,6 +14,7 @@ class RevokeEvent(BaseModel): requester: entities.slack.User user_account_assignment: sso.UserAccountAssignment permission_duration: timedelta + thread_ts: str | None = None class GroupRevokeEvent(BaseModel): @@ -22,6 +23,7 @@ class GroupRevokeEvent(BaseModel): requester: entities.slack.User group_assignment: sso.GroupAssignment permission_duration: timedelta + thread_ts: str | None = None class ScheduledGroupRevokeEvent(BaseModel): diff --git a/src/group.py b/src/group.py index ce28cd0..5c9d386 100644 --- a/src/group.py +++ b/src/group.py @@ -138,6 +138,7 @@ def handle_request_for_group_access_submittion( reason=request.reason, decision=decision, identity_store_id=identity_store_id, + thread_ts=slack_response["ts"], ) if decision.grant: @@ -247,6 +248,7 @@ def handle_group_button_click(body: dict, client: WebClient, context: BoltContex requester=requester, reason=payload.request.reason, identity_store_id=identity_store_id, + thread_ts=payload.thread_ts, ) cache_for_dublicate_requests.clear() if cfg.send_dm_if_user_not_in_channel and not is_user_in_channel: diff --git a/src/main.py b/src/main.py index 3240c39..39ed8b8 100644 --- a/src/main.py +++ b/src/main.py @@ -289,6 +289,7 @@ def handle_button_click(body: dict, client: WebClient, context: BoltContext) -> approver=approver, requester=requester, reason=payload.request.reason, + thread_ts=payload.thread_ts, ) cache_for_dublicate_requests.clear() if cfg.send_dm_if_user_not_in_channel and not is_user_in_channel: @@ -448,6 +449,7 @@ def handle_request_for_access_submittion( # noqa: PLR0915, PLR0912 approver=requester, requester=requester, reason=request.reason, + thread_ts=slack_response["ts"], ) if decision.grant: diff --git a/src/revoker.py b/src/revoker.py index 9846665..3dee375 100644 --- a/src/revoker.py +++ b/src/revoker.py @@ -182,6 +182,7 @@ def slack_notify_user_on_revoke( # noqa: PLR0913 sso_client: SSOAdminClient, identitystore_client: IdentityStoreClient, slack_client: slack_sdk.WebClient, + thread_ts: str | None = None, ) -> SlackResponse: mention = slack_helpers.create_slack_mention_by_principal_id( sso_user_id=( @@ -197,6 +198,7 @@ def slack_notify_user_on_revoke( # noqa: PLR0913 return slack_client.chat_postMessage( channel=cfg.slack_channel_id, text=f"Revoked role {permission_set.name} for user {mention} in account {account.name}", + thread_ts=thread_ts, ) @@ -206,6 +208,7 @@ def slack_notify_user_on_group_access_revoke( # noqa: PLR0913 sso_client: SSOAdminClient, identitystore_client: IdentityStoreClient, slack_client: slack_sdk.WebClient, + thread_ts: str | None = None, ) -> SlackResponse: mention = slack_helpers.create_slack_mention_by_principal_id( sso_user_id=group_assignment.user_principal_id, @@ -217,6 +220,7 @@ def slack_notify_user_on_group_access_revoke( # noqa: PLR0913 return slack_client.chat_postMessage( channel=cfg.slack_channel_id, text=f"User {mention} has been removed from the group {group_assignment.group_name}.", + thread_ts=thread_ts, ) @@ -270,6 +274,7 @@ def handle_scheduled_account_assignment_deletion( # noqa: PLR0913 sso_client=sso_client, identitystore_client=identitystore_client, slack_client=slack_client, + thread_ts=revoke_event.thread_ts, ) @@ -307,6 +312,7 @@ def handle_scheduled_group_assignment_deletion( # noqa: PLR0913 sso_client=sso_client, identitystore_client=identitystore_client, slack_client=slack_client, + thread_ts=group_revoke_event.thread_ts, ) diff --git a/src/schedule.py b/src/schedule.py index 6b2ffa3..01eb3ca 100644 --- a/src/schedule.py +++ b/src/schedule.py @@ -137,6 +137,7 @@ def schedule_revoke_event( approver: entities.slack.User, requester: entities.slack.User, user_account_assignment: sso.UserAccountAssignment, + thread_ts: str | None = None, ) -> scheduler_type_defs.CreateScheduleOutputTypeDef: logger.info("Scheduling revoke event") schedule_name = f"{cfg.revoker_function_name}" + datetime.now(timezone.utc).strftime("%Y-%m-%d-%H-%M-%S") @@ -147,6 +148,7 @@ def schedule_revoke_event( requester=requester, user_account_assignment=user_account_assignment, permission_duration=permission_duration, + thread_ts=thread_ts, ) logger.debug("Creating schedule", extra={"revoke_event": revoke_event}) return schedule_client.create_schedule( @@ -175,6 +177,7 @@ def schedule_group_revoke_event( approver: entities.slack.User, requester: entities.slack.User, group_assignment: sso.GroupAssignment, + thread_ts: str | None = None, ) -> scheduler_type_defs.CreateScheduleOutputTypeDef: logger.info("Scheduling revoke event") schedule_name = f"{cfg.revoker_function_name}" + datetime.now(timezone.utc).strftime("%Y-%m-%d-%H-%M-%S") @@ -184,6 +187,7 @@ def schedule_group_revoke_event( requester=requester, group_assignment=group_assignment, permission_duration=permission_duration, + thread_ts=thread_ts, ) get_and_delete_scheduled_revoke_event_if_already_exist(schedule_client, group_assignment) logger.debug("Creating schedule", extra={"revoke_event": revoke_event}) diff --git a/src/tests/test_model_serialization.py b/src/tests/test_model_serialization.py index bef1182..03af46b 100644 --- a/src/tests/test_model_serialization.py +++ b/src/tests/test_model_serialization.py @@ -1,6 +1,17 @@ """Tests for model serialization, particularly handling frozensets with nested Pydantic models.""" +from datetime import timedelta + +import entities +import sso from access_control import AccessRequestDecision, ApproveRequestDecision, DecisionReason +from events import ( + Event, + GroupRevokeEvent, + RevokeEvent, + ScheduledGroupRevokeEvent, + ScheduledRevokeEvent, +) from statement import Statement, GroupStatement @@ -174,3 +185,138 @@ def test_group_statement_dict_serialization(self): assert isinstance(result["approvers"], list) assert result["allow_self_approval"] is False assert result["approval_is_not_required"] is True + + +class TestRevokeEventSerialization: + """Test that RevokeEvent and GroupRevokeEvent serialize thread_ts correctly through EventBridge payload.""" + + def _sample_user(self): + return entities.slack.User(email="user@example.com", id="U123", real_name="Test User") + + def _sample_user_account_assignment(self): + return sso.UserAccountAssignment( + instance_arn="arn:aws:sso:::instance/ssoins-123", + account_id="123456789012", + permission_set_arn="arn:aws:sso:::permissionSet/ssoins-123/ps-123", + user_principal_id="user-principal-123", + ) + + def _sample_group_assignment(self): + return sso.GroupAssignment( + identity_store_id="d-123456789", + group_name="TestGroup", + group_id="group-123", + user_principal_id="user-principal-123", + membership_id="membership-123", + ) + + def test_revoke_event_with_thread_ts_json_roundtrip(self): + """Test RevokeEvent preserves thread_ts through JSON serialization (EventBridge payload).""" + event = RevokeEvent( + schedule_name="test-schedule", + approver=self._sample_user(), + requester=self._sample_user(), + user_account_assignment=self._sample_user_account_assignment(), + permission_duration=timedelta(hours=1), + thread_ts="1234567890.123456", + ) + + # Serialize to JSON (mimics EventBridge payload) + json_str = event.json() + + # Deserialize back + restored = RevokeEvent.model_validate_json(json_str) + + assert restored.thread_ts == "1234567890.123456" + + def test_revoke_event_without_thread_ts_backward_compat(self): + """Test RevokeEvent works without thread_ts (older scheduled jobs).""" + # JSON without thread_ts field (simulates older scheduled jobs) + event = RevokeEvent( + schedule_name="test-schedule", + approver=self._sample_user(), + requester=self._sample_user(), + user_account_assignment=self._sample_user_account_assignment(), + permission_duration=timedelta(hours=1), + ) + + json_str = event.json() + restored = RevokeEvent.model_validate_json(json_str) + + assert restored.thread_ts is None + + def test_group_revoke_event_with_thread_ts_json_roundtrip(self): + """Test GroupRevokeEvent preserves thread_ts through JSON serialization.""" + event = GroupRevokeEvent( + schedule_name="test-schedule", + approver=self._sample_user(), + requester=self._sample_user(), + group_assignment=self._sample_group_assignment(), + permission_duration=timedelta(hours=1), + thread_ts="1234567890.654321", + ) + + json_str = event.json() + restored = GroupRevokeEvent.model_validate_json(json_str) + + assert restored.thread_ts == "1234567890.654321" + + def test_group_revoke_event_without_thread_ts_backward_compat(self): + """Test GroupRevokeEvent works without thread_ts (older scheduled jobs).""" + event = GroupRevokeEvent( + schedule_name="test-schedule", + approver=self._sample_user(), + requester=self._sample_user(), + group_assignment=self._sample_group_assignment(), + permission_duration=timedelta(hours=1), + ) + + json_str = event.json() + restored = GroupRevokeEvent.model_validate_json(json_str) + + assert restored.thread_ts is None + + def test_scheduled_revoke_event_parses_thread_ts_from_nested_json(self): + """Test ScheduledRevokeEvent model_validator preserves thread_ts from JSON string.""" + revoke_event = RevokeEvent( + schedule_name="test-schedule", + approver=self._sample_user(), + requester=self._sample_user(), + user_account_assignment=self._sample_user_account_assignment(), + permission_duration=timedelta(hours=1), + thread_ts="1234567890.999999", + ) + + # This mimics the EventBridge payload structure where revoke_event is a JSON string + payload = { + "action": "event_bridge_revoke", + "revoke_event": revoke_event.json(), + } + + # Parse using Event (the root model used in revoker.py) + parsed = Event.model_validate(payload) + + assert isinstance(parsed.root, ScheduledRevokeEvent) + assert parsed.root.revoke_event.thread_ts == "1234567890.999999" + + def test_scheduled_group_revoke_event_parses_thread_ts_from_nested_json(self): + """Test ScheduledGroupRevokeEvent model_validator preserves thread_ts from JSON string.""" + group_revoke_event = GroupRevokeEvent( + schedule_name="test-schedule", + approver=self._sample_user(), + requester=self._sample_user(), + group_assignment=self._sample_group_assignment(), + permission_duration=timedelta(hours=1), + thread_ts="1234567890.111111", + ) + + # This mimics the EventBridge payload structure + payload = { + "action": "event_bridge_group_revoke", + "revoke_event": group_revoke_event.json(), + } + + parsed = Event.model_validate(payload) + + assert isinstance(parsed.root, ScheduledGroupRevokeEvent) + assert parsed.root.revoke_event.thread_ts == "1234567890.111111"