Merge pull request #15 from PostHog/tom/fix-2

Fix using permission set arn (again)

Author: Piccirello

src/access_control.py

@@ -43,9 +43,10 @@ def determine_affected_statements(
     account_id: str | None = None,
     permission_set_name: str | None = None,
     group_id: str | None = None,
+    permission_set_arn: str | None = None,
 ) -> FrozenSet[Statement] | FrozenSet[GroupStatement]:
     if isinstance(statements, FrozenSet) and all(isinstance(item, Statement) for item in statements):
-        return get_affected_statements(statements, account_id, permission_set_name)  # type: ignore # noqa: PGH003
+        return get_affected_statements(statements, account_id, permission_set_name, permission_set_arn)  # type: ignore # noqa: PGH003
 
     if isinstance(statements, FrozenSet) and all(isinstance(item, GroupStatement) for item in statements):
         return get_affected_group_statements(statements, group_id)  # type: ignore # noqa: PGH003
@@ -63,6 +64,7 @@ def make_decision_on_access_request(  # noqa: PLR0911, PLR0913
     account_id: str | None = None,
     group_id: str | None = None,
     user_group_ids: set[str] | None = None,
+    permission_set_arn: str | None = None,
 ) -> AccessRequestDecision:
     """Make a decision on an access request.
 
@@ -75,13 +77,14 @@ def make_decision_on_access_request(  # noqa: PLR0911, PLR0913
         user_group_ids: The set of SSO group IDs the user belongs to.
             If None, group-based filtering is skipped (backwards compatible).
             If an empty set, only statements without required_group_membership will match.
+        permission_set_arn: The ARN of the permission set (optional, for matching ARN-based config).
     """
     # Filter statements by user's group membership eligibility if user_group_ids provided
     # This is only applicable to Statement (not GroupStatement)
     if user_group_ids is not None and isinstance(statements, frozenset) and all(isinstance(s, Statement) for s in statements):
         statements = get_eligible_statements_for_user(statements, user_group_ids)  # type: ignore # noqa: PGH003
 
-    affected_statements = determine_affected_statements(statements, account_id, permission_set_name, group_id)
+    affected_statements = determine_affected_statements(statements, account_id, permission_set_name, group_id, permission_set_arn)
 
     decision_based_on_statements: set[Statement] | set[GroupStatement] = set()
     potential_approvers = set()
@@ -168,8 +171,9 @@ def make_decision_on_approve_request(  # noqa: PLR0913
     permission_set_name: str | None = None,
     account_id: str | None = None,
     group_id: str | None = None,
+    permission_set_arn: str | None = None,
 ) -> ApproveRequestDecision:
-    affected_statements = determine_affected_statements(statements, account_id, permission_set_name, group_id)
+    affected_statements = determine_affected_statements(statements, account_id, permission_set_name, group_id, permission_set_arn)
 
     for statement in affected_statements:
         if approver_email in statement.approvers:

src/main.py

@@ -258,6 +258,9 @@ def handle_button_click(body: dict, client: WebClient, context: BoltContext) ->
     cache_for_dublicate_requests["account_id"] = payload.request.account_id
     cache_for_dublicate_requests["permission_set_name"] = payload.request.permission_set_name
 
+    # Look up permission set to get ARN for matching and name for display
+    permission_set = sso.get_permission_set(sso_client, cfg.sso_instance_arn, payload.request.permission_set_name)
+
     if payload.action == entities.ApproverAction.Discard:
         blocks = slack_helpers.HeaderSectionBlock.set_status(
             blocks=payload.message["blocks"],
@@ -281,7 +284,7 @@ def handle_button_click(body: dict, client: WebClient, context: BoltContext) ->
             distinct_id=requester.email,
             properties={
                 "account_id": payload.request.account_id,
-                "permission_set": payload.request.permission_set_name,
+                "permission_set": permission_set.name,
                 "approver_email": approver.email,
                 "requester_email": requester.email,
             },
@@ -304,6 +307,7 @@ def handle_button_click(body: dict, client: WebClient, context: BoltContext) ->
         permission_set_name=payload.request.permission_set_name,
         approver_email=approver.email,
         requester_email=requester.email,
+        permission_set_arn=permission_set.arn,
     )
     logger.info("Decision on request was made", extra={"decision": decision.dict()})
 
@@ -349,7 +353,7 @@ def handle_button_click(body: dict, client: WebClient, context: BoltContext) ->
             distinct_id=requester.email,
             properties={
                 "account_id": payload.request.account_id,
-                "permission_set": payload.request.permission_set_name,
+                "permission_set": permission_set.name,
                 "approver_email": approver.email,
                 "requester_email": requester.email,
                 "duration_hours": payload.request.permission_duration.total_seconds() / 3600,
@@ -443,12 +447,16 @@ def handle_request_for_access_submittion(  # noqa: PLR0915, PLR0912
             user_principal_id=user_principal_id,
         )
 
+    # Look up permission set to get ARN for matching against ARN-based config
+    permission_set = sso.get_permission_set(sso_client, cfg.sso_instance_arn, request.permission_set_name)
+
     decision = access_control.make_decision_on_access_request(
         cfg.statements,
         account_id=request.account_id,
         permission_set_name=request.permission_set_name,
         requester_email=requester.email,
         user_group_ids=user_group_ids,
+        permission_set_arn=permission_set.arn,
     )
     logger.info("Decision on request was made", extra={"decision": decision.dict()})
 
@@ -457,7 +465,7 @@ def handle_request_for_access_submittion(  # noqa: PLR0915, PLR0912
         distinct_id=requester.email,
         properties={
             "account_id": request.account_id,
-            "permission_set": request.permission_set_name,
+            "permission_set": permission_set.name,
             "requester_email": requester.email,
             "decision_reason": decision.reason.value,
             "granted": decision.grant,
@@ -478,7 +486,7 @@ def handle_request_for_access_submittion(  # noqa: PLR0915, PLR0912
             slack_client=client,
             requester_slack_id=request.requester_slack_id,
             account=account,
-            role_name=request.permission_set_name,
+            role_name=permission_set.name,
             reason=request.reason,
             permission_duration=request.permission_duration,
             show_buttons=show_buttons,
@@ -590,7 +598,7 @@ def handle_request_for_access_submittion(  # noqa: PLR0915, PLR0912
             distinct_id=requester.email,
             properties={
                 "account_id": request.account_id,
-                "permission_set": request.permission_set_name,
+                "permission_set": permission_set.name,
                 "approver_email": requester.email,
                 "requester_email": requester.email,
                 "duration_hours": request.permission_duration.total_seconds() / 3600,

src/revoker.py

@@ -463,9 +463,7 @@ def slack_notify_user_on_revoke(  # noqa: PLR0913
             )
         # Delete the early revoke button from the thread
         slack_helpers.delete_early_revoke_button(slack_client, cfg.slack_channel_id, thread_ts)
-        # Simplified message for threads (context already in thread)
-        # Change the "Access revoked." message to something friendly and maybe even a bit cheeky. It's too serious. At a minimum "Access expired""ended"/"completed"
-        text = "Access revoked."
+        text = "Session complete."
     else:
         # Full message when not in a thread
         mention = slack_helpers.create_slack_mention_by_principal_id(
@@ -515,9 +513,7 @@ def slack_notify_user_on_group_access_revoke(  # noqa: PLR0913
             )
         # Delete the early revoke button from the thread
         slack_helpers.delete_early_revoke_button(slack_client, cfg.slack_channel_id, thread_ts)
-        # Simplified message for threads (context already in thread)
-        # Change the "Access revoked." message to something friendly and maybe even a bit cheeky. It's too serious. At a minimum "Access expired""ended"/"completed"
-        text = "Access revoked."
+        text = "Session complete."
     else:
         # Full message when not in a thread
         mention = slack_helpers.create_slack_mention_by_principal_id(

src/slack_helpers.py

@@ -137,7 +137,7 @@ def build_select_permission_set_input_block(cls, permission_sets: list[entities.
                 action_id=cls.PERMISSION_SET_ACTION_ID,
                 placeholder=PlainTextObject(text="Select permission set"),
                 options=[
-                    Option(text=PlainTextObject(text=permission_set.name), value=permission_set.name)
+                    Option(text=PlainTextObject(text=permission_set.name), value=permission_set.arn)
                     for permission_set in sorted_permission_sets
                 ],
             ),

src/sso.py

@@ -500,7 +500,7 @@ def get_permission_sets_from_config_with_cache(
     if "*" in cfg.permission_sets:
         return all_permission_sets
     else:
-        return [ps for ps in all_permission_sets if ps.name in cfg.permission_sets]
+        return [ps for ps in all_permission_sets if ps.name in cfg.permission_sets or ps.arn in cfg.permission_sets]
 
 
 def get_account_assignment_information(

src/statement.py

@@ -32,14 +32,26 @@ class Statement(BaseStatement):
     resource_type: ResourceType = Field(default=ResourceType.Account, frozen=True)
     resource: FrozenSet[Union[AWSAccountId, WildCard]]
 
-    def affects(self, account_id: str, permission_set_name: str) -> bool:  # noqa: ANN101
-        return (account_id in self.resource or "*" in self.resource) and (
-            permission_set_name in self.permission_set or "*" in self.permission_set
+    def affects(self, account_id: str, permission_set_name: str, permission_set_arn: str | None = None) -> bool:  # noqa: ANN101
+        account_match = account_id in self.resource or "*" in self.resource
+        ps_match = (
+            permission_set_name in self.permission_set
+            or "*" in self.permission_set
+            or (permission_set_arn is not None and permission_set_arn in self.permission_set)
         )
-
-
-def get_affected_statements(statements: FrozenSet[Statement], account_id: str, permission_set_name: str) -> FrozenSet[Statement]:
-    return frozenset(statement for statement in statements if statement.affects(account_id, permission_set_name))
+        return account_match and ps_match
+
+
+def get_affected_statements(
+    statements: FrozenSet[Statement],
+    account_id: str,
+    permission_set_name: str,
+    permission_set_arn: str | None = None,
+) -> FrozenSet[Statement]:
+    return frozenset(
+        statement for statement in statements
+        if statement.affects(account_id, permission_set_name, permission_set_arn)
+    )
 
 
 def get_permission_sets_for_account(statements: FrozenSet[Statement], account_id: str) -> set[str]: