feat: add get_credentials_by method for OAuth credential exchange with integration filtering

Author: mconflitti-pbc

src/posit/connect/oauth/oauth.py

@@ -1,15 +1,20 @@
 from __future__ import annotations
 
 import os
+import re
 from enum import Enum
 
 from typing_extensions import TYPE_CHECKING, Optional, TypedDict
 
+from posit.connect.oauth.associations import ContentItemAssociations
+
 from ..resources import Resources
 from .integrations import Integrations
 from .sessions import Sessions
 
 if TYPE_CHECKING:
+    from posit.connect.external import aws
+
     from ..context import Context
 
 GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
@@ -23,6 +28,25 @@ class OAuthTokenType(str, Enum):
     USER_SESSION_TOKEN = "urn:posit:connect:user-session-token"
 
 
+class OAuthIntegrationAuthType(str, Enum):
+    """OAuth integration authentication type."""
+
+    VIEWER = "Viewer"
+    SERVICE_ACCOUNT = "Service Account"
+    VISITOR_API_KEY = "Visitor API Key"
+
+
+class OAuthIntegrationType(str, Enum):
+    """OAuth integration type."""
+
+    AWS = "aws"
+    AZURE = "azure"
+    CONNECT = "connect"
+    SNOWFLAKE = "snowflake"
+    CUSTOM = "custom"
+    # TODO add the rest
+
+
 def _get_content_session_token() -> str:
     """Return the content session token.
 
@@ -99,6 +123,74 @@ def get_content_credentials(
         response = self._ctx.client.post(self._path, data=data)
         return Credentials(**response.json())
 
+    def get_credentials_by(
+        self,
+        user_session_token: str,
+        content_session_token: Optional[str] = None,
+        integration_type: Optional[OAuthIntegrationType] = None,
+        auth_type: Optional[OAuthIntegrationAuthType] = None,
+        name: Optional[str | re.Pattern] = None,
+        guid: Optional[str] = None,
+    ) -> Credentials | aws.Credentials:
+        """Perform an oauth credential exchange for all integrations associated with the current content item."""
+        content_guid = os.getenv("CONNECT_CONTENT_GUID")
+        if not content_guid:
+            raise ValueError("CONNECT_CONTENT_GUID environment variable is required.")
+        content_associations = ContentItemAssociations(self._ctx, content_guid=content_guid).find()
+        # associations format: [{
+        #     content_guid: uuid
+        #     app_guid: uuid
+        #     oauth_integration_guid: uuid
+        #     oauth_integration_name: string
+        #     oauth_integration_description: string┃null
+        #     oauth_integration_template: string
+        #     oauth_integration_auth_type: string
+        #     created_time: date-time
+        # }]
+        for association in content_associations:
+            match = True
+
+            if (
+                integration_type is not None
+                and association.get("oauth_integration_template") != integration_type
+            ):
+                match = False
+
+            if (
+                auth_type is not None
+                and association.get("oauth_integration_auth_type") != auth_type
+            ):
+                match = False
+
+            if name is not None:
+                integration_name = association.get("oauth_integration_name", "")
+                if isinstance(name, re.Pattern):
+                    if not name.search(integration_name):
+                        match = False
+                else:
+                    if integration_name != name:
+                        match = False
+
+            if guid is not None and association.get("oauth_integration_guid") != guid:
+                match = False
+
+            if match:
+                # Use the first matching association to get credentials
+                if association.get("oauth_integration_auth_type") in [
+                    OAuthIntegrationAuthType.VIEWER,
+                    OAuthIntegrationAuthType.VISITOR_API_KEY,
+                ]:
+                    return self.get_credentials(
+                        user_session_token=user_session_token,
+                        audience=association.get("oauth_integration_guid"),
+                    )
+                return self.get_content_credentials(
+                    content_session_token=content_session_token,
+                    audience=association.get("oauth_integration_guid"),
+                )
+
+        raise ValueError("No matching OAuth integration found for the specified criteria.")
+
 
 class Credentials(TypedDict, total=False):
     access_token: str