Add custom org roles endpoints (#301)

* Add custom org roles endpoints

* use org policy for authorization checks

* ignore type check

* reformat

---------

Co-authored-by: Stytch Codegen Bot <[email protected]>
Co-authored-by: Evelyn Taylor-McGregor <[email protected]>

Author: 3 people

stytch/b2b/api/rbac.py

@@ -8,6 +8,7 @@
 
 from typing import Any, Dict
 
+from stytch.b2b.api.rbac_organizations import Organizations
 from stytch.b2b.models.rbac import PolicyResponse
 from stytch.core.api_base import ApiBase
 from stytch.core.http.client import AsyncClient, SyncClient
@@ -20,6 +21,11 @@ def __init__(
         self.api_base = api_base
         self.sync_client = sync_client
         self.async_client = async_client
+        self.organizations = Organizations(
+            api_base=self.api_base,
+            sync_client=self.sync_client,
+            async_client=self.async_client,
+        )
 
     def policy(
         self,

stytch/b2b/api/rbac_organizations.py

@@ -0,0 +1,122 @@
+# !!!
+# WARNING: This file is autogenerated
+# Only modify code within MANUAL() sections
+# or your changes may be overwritten later!
+# !!!
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Union
+
+from stytch.b2b.models.rbac import OrgPolicy
+from stytch.b2b.models.rbac_organizations import (
+    GetOrgPolicyResponse,
+    SetOrgPolicyResponse,
+)
+from stytch.core.api_base import ApiBase
+from stytch.core.http.client import AsyncClient, SyncClient
+
+
+class Organizations:
+    def __init__(
+        self, api_base: ApiBase, sync_client: SyncClient, async_client: AsyncClient
+    ) -> None:
+        self.api_base = api_base
+        self.sync_client = sync_client
+        self.async_client = async_client
+
+    def get_org_policy(
+        self,
+        organization_id: str,
+    ) -> GetOrgPolicyResponse:
+        """Get the active RBAC Policy for a specific Organization within your Stytch Project. An Organization RBAC Policy contains the roles that have been defined specifically for that organization, allowing for organization-specific permissioning models.
+
+        This endpoint returns the organization-scoped roles that supplement the project-level RBAC policy. Organization policies allow you to define custom roles that are specific to individual organizations within your project.
+
+        When using the backend SDKs, the RBAC Policy will be cached to allow for local evaluations, eliminating the need for an extra request to Stytch. The policy will be refreshed if an authorization check is requested and the RBAC policy was last updated more than 5 minutes ago.
+
+        Organization-specific roles can be created and managed through this API endpoint, providing fine-grained control over permissions at the organization level.
+
+        Check out the [RBAC overview](https://stytch.com/docs/b2b/guides/rbac/overview) to learn more about Stytch's RBAC permissioning model and organization-scoped policies.
+
+        Fields:
+          - organization_id: Globally unique UUID that identifies a specific Organization. The `organization_id` is critical to perform operations on an Organization, so be sure to preserve this value. You may also use the organization_slug or organization_external_id here as a convenience.
+        """  # noqa
+        headers: Dict[str, str] = {}
+        data: Dict[str, Any] = {
+            "organization_id": organization_id,
+        }
+
+        url = self.api_base.url_for(
+            "/v1/b2b/rbac/organizations/{organization_id}", data
+        )
+        res = self.sync_client.get(url, data, headers)
+        return GetOrgPolicyResponse.from_json(res.response.status_code, res.json)
+
+    async def get_org_policy_async(
+        self,
+        organization_id: str,
+    ) -> GetOrgPolicyResponse:
+        """Get the active RBAC Policy for a specific Organization within your Stytch Project. An Organization RBAC Policy contains the roles that have been defined specifically for that organization, allowing for organization-specific permissioning models.
+
+        This endpoint returns the organization-scoped roles that supplement the project-level RBAC policy. Organization policies allow you to define custom roles that are specific to individual organizations within your project.
+
+        When using the backend SDKs, the RBAC Policy will be cached to allow for local evaluations, eliminating the need for an extra request to Stytch. The policy will be refreshed if an authorization check is requested and the RBAC policy was last updated more than 5 minutes ago.
+
+        Organization-specific roles can be created and managed through this API endpoint, providing fine-grained control over permissions at the organization level.
+
+        Check out the [RBAC overview](https://stytch.com/docs/b2b/guides/rbac/overview) to learn more about Stytch's RBAC permissioning model and organization-scoped policies.
+
+        Fields:
+          - organization_id: Globally unique UUID that identifies a specific Organization. The `organization_id` is critical to perform operations on an Organization, so be sure to preserve this value. You may also use the organization_slug or organization_external_id here as a convenience.
+        """  # noqa
+        headers: Dict[str, str] = {}
+        data: Dict[str, Any] = {
+            "organization_id": organization_id,
+        }
+
+        url = self.api_base.url_for(
+            "/v1/b2b/rbac/organizations/{organization_id}", data
+        )
+        res = await self.async_client.get(url, data, headers)
+        return GetOrgPolicyResponse.from_json(res.response.status, res.json)
+
+    def set_org_policy(
+        self,
+        organization_id: str,
+        org_policy: Optional[Union[OrgPolicy, Dict[str, Any]]] = None,
+    ) -> SetOrgPolicyResponse:
+        headers: Dict[str, str] = {}
+        data: Dict[str, Any] = {
+            "organization_id": organization_id,
+        }
+        if org_policy is not None:
+            data["org_policy"] = (
+                org_policy if isinstance(org_policy, dict) else org_policy.dict()
+            )
+
+        url = self.api_base.url_for(
+            "/v1/b2b/rbac/organizations/{organization_id}", data
+        )
+        res = self.sync_client.put(url, data, headers)
+        return SetOrgPolicyResponse.from_json(res.response.status_code, res.json)
+
+    async def set_org_policy_async(
+        self,
+        organization_id: str,
+        org_policy: Optional[OrgPolicy] = None,
+    ) -> SetOrgPolicyResponse:
+        headers: Dict[str, str] = {}
+        data: Dict[str, Any] = {
+            "organization_id": organization_id,
+        }
+        if org_policy is not None:
+            data["org_policy"] = (
+                org_policy if isinstance(org_policy, dict) else org_policy.dict()
+            )
+
+        url = self.api_base.url_for(
+            "/v1/b2b/rbac/organizations/{organization_id}", data
+        )
+        res = await self.async_client.put(url, data, headers)
+        return SetOrgPolicyResponse.from_json(res.response.status, res.json)

stytch/b2b/api/sessions.py

@@ -990,7 +990,9 @@ def authenticate_jwt_local(
                 raise ValueError("Invalid roles claim. Expected a list of strings.")
 
             rbac_local.perform_authorization_check(
-                policy=self.policy_cache.get(),
+                policy=self.policy_cache.get_with_org(
+                    local_resp.member_session.organization_id
+                ),
                 subject_roles=local_resp.roles_claim,
                 subject_org_id=local_resp.member_session.organization_id,
                 authorization_check=authorization_check,
@@ -1019,7 +1021,9 @@ async def authenticate_jwt_local_async(
                 raise ValueError("Invalid roles claim. Expected a list of strings.")
 
             rbac_local.perform_authorization_check(
-                policy=await self.policy_cache.get_async(),
+                policy=await self.policy_cache.get_with_org_async(
+                    local_resp.member_session.organization_id
+                ),
                 subject_roles=local_resp.roles_claim,
                 subject_org_id=local_resp.member_session.organization_id,
                 authorization_check=authorization_check,

stytch/b2b/models/organizations.py

@@ -90,6 +90,17 @@ def add_headers(self, headers: Dict[str, str]) -> Dict[str, str]:
         return headers
 
 
+class CustomRolePermission(pydantic.BaseModel):
+    resource_id: str
+    actions: List[str]
+
+
+class CustomRole(pydantic.BaseModel):
+    role_id: str
+    description: str
+    permissions: List[CustomRolePermission]
+
+
 class DeleteRequestOptions(pydantic.BaseModel):
     """
     Fields:
@@ -400,6 +411,7 @@ class Organization(pydantic.BaseModel):
       `NOT_ALLOWED` – no third party Connected Apps are permitted.
 
       - allowed_third_party_connected_apps: An array of third party Connected App IDs that are allowed for the Organization. Only used when the Organization's `third_party_connected_apps_allowed_type` is `RESTRICTED`.
+      - custom_roles: (no documentation yet)
       - trusted_metadata: An arbitrary JSON object for storing application-specific data or identity-provider-specific data.
       - created_at: The timestamp of the Organization's creation. Values conform to the RFC 3339 standard and are expressed in UTC, e.g. `2021-12-29T12:33:09Z`.
       - updated_at: The timestamp of when the Organization was last updated. Values conform to the RFC 3339 standard and are expressed in UTC, e.g. `2021-12-29T12:33:09Z`.
@@ -431,6 +443,7 @@ class Organization(pydantic.BaseModel):
     allowed_first_party_connected_apps: List[str]
     third_party_connected_apps_allowed_type: str
     allowed_third_party_connected_apps: List[str]
+    custom_roles: List[CustomRole]
     trusted_metadata: Optional[Dict[str, Any]] = None
     created_at: Optional[datetime.datetime] = None
     updated_at: Optional[datetime.datetime] = None

stytch/b2b/models/rbac.py

@@ -122,6 +122,15 @@ class PolicyRole(pydantic.BaseModel):
     permissions: List[PolicyRolePermission]
 
 
+class OrgPolicy(pydantic.BaseModel):
+    """
+    Fields:
+      - roles: An array of [Role objects](https://stytch.com/docs/b2b/api/rbac-role-object).
+    """  # noqa
+
+    roles: List[PolicyRole]
+
+
 class PolicyScopePermission(pydantic.BaseModel):
     resource_id: str
     actions: List[str]

stytch/b2b/models/rbac_organizations.py

@@ -0,0 +1,25 @@
+# !!!
+# WARNING: This file is autogenerated
+# Only modify code within MANUAL() sections
+# or your changes may be overwritten later!
+# !!!
+
+from __future__ import annotations
+
+from typing import Optional
+
+from stytch.b2b.models.rbac import OrgPolicy
+from stytch.core.response_base import ResponseBase
+
+
+class GetOrgPolicyResponse(ResponseBase):
+    """Response type for `Organizations.get_org_policy`.
+    Fields:
+      - org_policy: The organization-specific RBAC Policy that contains roles defined for this organization. Organization policies supplement the project-level RBAC policy with additional roles that are specific to the organization.
+    """  # noqa
+
+    org_policy: Optional[OrgPolicy] = None
+
+
+class SetOrgPolicyResponse(ResponseBase):
+    org_policy: Optional[OrgPolicy] = None

stytch/consumer/api/otp.py

@@ -10,7 +10,7 @@
 
 from stytch.consumer.api.otp_email import Email
 from stytch.consumer.api.otp_sms import Sms
-from stytch.consumer.api.otp_whatsapp import Whatsapp
+from stytch.consumer.api.otp_whatsapp import WhatsApp
 from stytch.consumer.models.attribute import Attributes
 from stytch.consumer.models.magic_links import Options
 from stytch.consumer.models.otp import AuthenticateResponse
@@ -30,7 +30,7 @@ def __init__(
             sync_client=self.sync_client,
             async_client=self.async_client,
         )
-        self.whatsapp = Whatsapp(
+        self.whatsapp = WhatsApp(
             api_base=self.api_base,
             sync_client=self.sync_client,
             async_client=self.async_client,

stytch/consumer/api/otp_whatsapp.py

@@ -19,7 +19,7 @@
 from stytch.core.http.client import AsyncClient, SyncClient
 
 
-class Whatsapp:
+class WhatsApp:
     def __init__(
         self, api_base: ApiBase, sync_client: SyncClient, async_client: AsyncClient
     ) -> None:

stytch/consumer/models/otp_whatsapp.py

@@ -34,7 +34,7 @@ class SendRequestLocale(str, enum.Enum):
 
 
 class LoginOrCreateResponse(ResponseBase):
-    """Response type for `Whatsapp.login_or_create`.
+    """Response type for `WhatsApp.login_or_create`.
     Fields:
       - user_id: The unique ID of the affected User.
       - phone_id: The unique ID for the phone number.
@@ -47,7 +47,7 @@ class LoginOrCreateResponse(ResponseBase):
 
 
 class SendResponse(ResponseBase):
-    """Response type for `Whatsapp.send`.
+    """Response type for `WhatsApp.send`.
     Fields:
       - user_id: The unique ID of the affected User.
       - phone_id: The unique ID for the phone number.

stytch/shared/policy_cache.py

@@ -1,19 +1,45 @@
-from typing import Optional
+import time
+from typing import Dict, Optional
 
 from stytch.b2b.api.rbac import RBAC
-from stytch.b2b.models.rbac import Policy
+from stytch.b2b.models.rbac import OrgPolicy, Policy
 from stytch.shared.lazy_cache import LazyCache
 
 DEFAULT_REFRESH_INTERVAL = 10 * 60  # 10 minutes
 
 
+def _merge_policies(project_policy: Policy, org_policy: Optional[OrgPolicy]) -> Policy:
+    """Merge project-level and organization-level policies.
+
+    Organization policies supplement the project policy by adding additional roles.
+    The merged policy combines roles from both, with org roles added to project roles.
+    Resources and scopes come only from the project policy.
+    """
+    if org_policy is None:
+        return project_policy
+
+    return Policy(
+        roles=list(project_policy.roles) + list(org_policy.roles),
+        resources=project_policy.resources,
+        scopes=project_policy.scopes,
+    )
+
+
+class _OrgPolicyCacheEntry:
+    def __init__(self, org_policy: Optional[OrgPolicy], last_refresh_time: float):
+        self.org_policy = org_policy
+        self.last_refresh_time = last_refresh_time
+
+
 class PolicyCache(LazyCache[Policy]):
     def __init__(
         self,
         rbac: RBAC,
         refresh_interval_seconds: int = DEFAULT_REFRESH_INTERVAL,
     ) -> None:
         self.rbac = rbac
+        self.refresh_interval_seconds = refresh_interval_seconds
+        self._org_cache: Dict[str, _OrgPolicyCacheEntry] = {}
         super().__init__(
             reload_func=self.reload_policy,
             async_reload_func=self.reload_policy_async,
@@ -29,3 +55,49 @@ async def reload_policy_async(self, _: Optional[Policy]) -> Policy:
         resp = await self.rbac.policy_async()
         assert resp.policy is not None
         return resp.policy
+
+    def _org_needs_refresh(self, organization_id: str) -> bool:
+        if organization_id not in self._org_cache:
+            return True
+        entry = self._org_cache[organization_id]
+        return time.time() - entry.last_refresh_time > self.refresh_interval_seconds
+
+    def _get_org_policy(self, organization_id: str) -> Optional[OrgPolicy]:
+        """Get the organization policy, fetching and caching if needed."""
+        if self._org_needs_refresh(organization_id):
+            resp = self.rbac.organizations.get_org_policy(organization_id)
+            self._org_cache[organization_id] = _OrgPolicyCacheEntry(
+                org_policy=resp.org_policy,
+                last_refresh_time=time.time(),
+            )
+        return self._org_cache[organization_id].org_policy
+
+    async def _get_org_policy_async(self, organization_id: str) -> Optional[OrgPolicy]:
+        """Get the organization policy, fetching and caching if needed (async)."""
+        if self._org_needs_refresh(organization_id):
+            resp = await self.rbac.organizations.get_org_policy_async(organization_id)
+            self._org_cache[organization_id] = _OrgPolicyCacheEntry(
+                org_policy=resp.org_policy,
+                last_refresh_time=time.time(),
+            )
+        return self._org_cache[organization_id].org_policy
+
+    def get_with_org(self, organization_id: str) -> Policy:
+        """Get the project policy merged with the organization policy.
+
+        This fetches both the project-level RBAC policy and the organization-specific
+        policy, merging the org roles into the project policy for authorization checks.
+        """
+        project_policy = self.get()
+        org_policy = self._get_org_policy(organization_id)
+        return _merge_policies(project_policy, org_policy)
+
+    async def get_with_org_async(self, organization_id: str) -> Policy:
+        """Get the project policy merged with the organization policy (async).
+
+        This fetches both the project-level RBAC policy and the organization-specific
+        policy, merging the org roles into the project policy for authorization checks.
+        """
+        project_policy = await self.get_async()
+        org_policy = await self._get_org_policy_async(organization_id)
+        return _merge_policies(project_policy, org_policy)