Fix #400 token rotation feature support

slackapi · Jul 15, 2021 · 6bfc7bb · 6bfc7bb
1 parent 1b4eefc
commit 6bfc7bb
Show file tree

Hide file tree

Showing 12 changed files with 171 additions and 2 deletions.
diff --git a/examples/oauth_sqlite3_app.py b/examples/oauth_sqlite3_app.py
@@ -5,7 +5,12 @@
 from slack_bolt import App
 from slack_bolt.oauth import OAuthFlow
 
-app = App(oauth_flow=OAuthFlow.sqlite3(database="./slackapp.db"))
+app = App(
+    oauth_flow=OAuthFlow.sqlite3(
+        database="./slackapp.db",
+        token_rotation_expiration_minutes=60 * 24,  # for testing
+    )
+)
 
 
 @app.event("app_mention")
@@ -14,6 +19,12 @@ def handle_app_mentions(body, say, logger):
     say("What's up?")
 
 
+@app.command("/token-rotation-modal")
+def handle_some_command(ack, body, logger):
+    ack()
+    logger.info(body)
+
+
 if __name__ == "__main__":
     app.start(3000)
 

diff --git a/examples/oauth_sqlite3_app_bot_only.py b/examples/oauth_sqlite3_app_bot_only.py
@@ -0,0 +1,38 @@
+import logging
+
+logging.basicConfig(level=logging.DEBUG)
+
+from slack_bolt import App
+from slack_bolt.oauth import OAuthFlow
+
+app = App(
+    oauth_flow=OAuthFlow.sqlite3(
+        database="./slackapp.db",
+        token_rotation_expiration_minutes=60 * 24,  # for testing
+    ),
+    installation_store_bot_only=True,
+)
+
+
+@app.event("app_mention")
+def handle_app_mentions(body, say, logger):
+    logger.info(body)
+    say("What's up?")
+
+
+@app.command("/token-rotation-modal")
+def handle_some_command(ack, body, logger):
+    ack()
+    logger.info(body)
+
+
+if __name__ == "__main__":
+    app.start(3000)
+
+# pip install slack_bolt
+# export SLACK_SIGNING_SECRET=***
+# export SLACK_BOT_TOKEN=xoxb-***
+# export SLACK_CLIENT_ID=111.111
+# export SLACK_CLIENT_SECRET=***
+# export SLACK_SCOPES=app_mentions:read,channels:history,im:history,chat:write
+# python oauth_app.py
diff --git a/setup.py b/setup.py
@@ -41,7 +41,7 @@
     ),
     include_package_data=True,  # MANIFEST.in
     install_requires=[
-        "slack_sdk>=3.5.0,<4",
+        "slack_sdk>=3.8.0rc2,<4",
     ],
     setup_requires=["pytest-runner==5.2"],
     tests_require=test_dependencies,

diff --git a/slack_bolt/adapter/aws_lambda/lambda_s3_oauth_flow.py b/slack_bolt/adapter/aws_lambda/lambda_s3_oauth_flow.py
@@ -63,6 +63,8 @@ def __init__(
         # In this case, the /slack/events endpoint doesn't work along with the OAuth flow.
         settings.authorize = InstallationStoreAuthorize(
             logger=logger,
+            client_id=settings.client_id,
+            client_secret=settings.client_secret,
             installation_store=settings.installation_store,
             bot_only=settings.installation_store_bot_only,
         )

diff --git a/slack_bolt/app/app.py b/slack_bolt/app/app.py
@@ -218,8 +218,11 @@ def message_hello(message, say):
 
         self._installation_store: Optional[InstallationStore] = installation_store
         if self._installation_store is not None and self._authorize is None:
+            settings = oauth_flow.settings if oauth_flow is not None else oauth_settings
             self._authorize = InstallationStoreAuthorize(
                 installation_store=self._installation_store,
+                client_id=settings.client_id if settings is not None else None,
+                client_secret=settings.client_secret if settings is not None else None,
                 logger=self._framework_logger,
                 bot_only=installation_store_bot_only,
             )

diff --git a/slack_bolt/app/async_app.py b/slack_bolt/app/async_app.py
@@ -229,8 +229,11 @@ async def message_hello(message, say):  # async function
             AsyncInstallationStore
         ] = installation_store
         if self._async_installation_store is not None and self._async_authorize is None:
+            settings = oauth_flow.settings if oauth_flow is not None else oauth_settings
             self._async_authorize = AsyncInstallationStoreAuthorize(
                 installation_store=self._async_installation_store,
+                client_id=settings.client_id if settings is not None else None,
+                client_secret=settings.client_secret if settings is not None else None,
                 logger=self._framework_logger,
                 bot_only=installation_store_bot_only,
             )

diff --git a/slack_bolt/authorization/async_authorize.py b/slack_bolt/authorization/async_authorize.py
@@ -7,10 +7,12 @@
 from slack_sdk.oauth.installation_store.async_installation_store import (
     AsyncInstallationStore,
 )
+from slack_sdk.oauth.token_rotation.async_rotator import AsyncTokenRotator
 
 from slack_bolt.authorization.async_authorize_args import AsyncAuthorizeArgs
 from slack_bolt.authorization import AuthorizeResult
 from slack_bolt.context.async_context import AsyncBoltContext
+from slack_bolt.error import BoltError
 
 
 class AsyncAuthorize:
@@ -94,12 +96,18 @@ class AsyncInstallationStoreAuthorize(AsyncAuthorize):
     authorize_result_cache: Dict[str, AuthorizeResult]
     find_installation_available: Optional[bool]
     find_bot_available: Optional[bool]
+    token_rotator: Optional[AsyncTokenRotator]
+
+    _config_error_message: str = "AsyncInstallationStore with client_id/client_secret are required for token rotation"
 
     def __init__(
         self,
         *,
         logger: Logger,
         installation_store: AsyncInstallationStore,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        token_rotation_expiration_minutes: Optional[int] = None,
         # For v1.0.x compatibility and people who still want its simplicity
         # use only InstallationStore#find_bot(enterprise_id, team_id)
         bot_only: bool = False,
@@ -112,6 +120,16 @@ def __init__(
         self.authorize_result_cache = {}
         self.find_installation_available = None
         self.find_bot_available = None
+        if client_id is not None and client_secret is not None:
+            self.token_rotator = AsyncTokenRotator(
+                client_id=client_id,
+                client_secret=client_secret,
+            )
+        else:
+            self.token_rotator = None
+        self.token_rotation_expiration_minutes = (
+            token_rotation_expiration_minutes or 120
+        )
 
     async def __call__(
         self,
@@ -171,6 +189,20 @@ async def __call__(
                         installation.user_token,
                     )
 
+                    if installation.user_refresh_token is not None:
+                        if self.token_rotator is None:
+                            raise BoltError(self._config_error_message)
+                        refreshed = await self.token_rotator.perform_token_rotation(
+                            installation=installation,
+                            token_rotation_expiration_minutes=self.token_rotation_expiration_minutes,
+                        )
+                        if refreshed is not None:
+                            await self.installation_store.async_save(refreshed)
+                            bot_token, user_token = (
+                                refreshed.bot_token,
+                                refreshed.user_token,
+                            )
+
             except NotImplementedError as _:
                 self.find_installation_available = False
 
@@ -199,6 +231,18 @@ async def __call__(
             except Exception as e:
                 self.logger.info(f"Failed to call find_bot method: {e}")
 
+            if bot.bot_refresh_token is not None:
+                # Token rotation
+                if self.token_rotator is None:
+                    raise BoltError(self._config_error_message)
+                refreshed = await self.token_rotator.perform_bot_token_rotation(
+                    bot=bot,
+                    token_rotation_expiration_minutes=self.token_rotation_expiration_minutes,
+                )
+                if refreshed is not None:
+                    await self.installation_store.async_save_bot(refreshed)
+                    bot_token = refreshed.bot_token
+
         token: Optional[str] = bot_token or user_token
         if token is None:
             # No valid token was found

diff --git a/slack_bolt/authorization/authorize.py b/slack_bolt/authorization/authorize.py
@@ -1,15 +1,18 @@
 import inspect
+import os
 from logging import Logger
 from typing import Optional, Callable, Dict, Any
 
 from slack_sdk.errors import SlackApiError
 from slack_sdk.oauth import InstallationStore
 from slack_sdk.oauth.installation_store import Bot
 from slack_sdk.oauth.installation_store.models.installation import Installation
+from slack_sdk.oauth.token_rotation.rotator import TokenRotator
 
 from slack_bolt.authorization.authorize_args import AuthorizeArgs
 from slack_bolt.authorization.authorize_result import AuthorizeResult
 from slack_bolt.context.context import BoltContext
+from slack_bolt.error import BoltError
 
 
 class Authorize:
@@ -97,12 +100,20 @@ class InstallationStoreAuthorize(Authorize):
     bot_only: bool
     find_installation_available: bool
     find_bot_available: bool
+    token_rotator: Optional[TokenRotator]
+
+    _config_error_message: str = (
+        "InstallationStore with client_id/client_secret are required for token rotation"
+    )
 
     def __init__(
         self,
         *,
         logger: Logger,
         installation_store: InstallationStore,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        token_rotation_expiration_minutes: Optional[int] = None,
         # For v1.0.x compatibility and people who still want its simplicity
         # use only InstallationStore#find_bot(enterprise_id, team_id)
         bot_only: bool = False,
@@ -117,6 +128,16 @@ def __init__(
             installation_store, "find_installation"
         )
         self.find_bot_available = hasattr(installation_store, "find_bot")
+        if client_id is not None and client_secret is not None:
+            self.token_rotator = TokenRotator(
+                client_id=client_id,
+                client_secret=client_secret,
+            )
+        else:
+            self.token_rotator = None
+        self.token_rotation_expiration_minutes = (
+            token_rotation_expiration_minutes or 120
+        )
 
     def __call__(
         self,
@@ -165,6 +186,19 @@ def __call__(
                         installation.bot_token,
                         installation.user_token,
                     )
+                    if installation.user_refresh_token is not None:
+                        if self.token_rotator is None:
+                            raise BoltError(self._config_error_message)
+                        refreshed = self.token_rotator.perform_token_rotation(
+                            installation=installation,
+                            minutes_before_expiration=self.token_rotation_expiration_minutes,
+                        )
+                        if refreshed is not None:
+                            self.installation_store.save(refreshed)
+                            bot_token, user_token = (
+                                refreshed.bot_token,
+                                refreshed.user_token,
+                            )
 
             except NotImplementedError as _:
                 self.find_installation_available = False
@@ -194,6 +228,18 @@ def __call__(
             except Exception as e:
                 self.logger.info(f"Failed to call find_bot method: {e}")
 
+            if bot.bot_refresh_token is not None:
+                # Token rotation
+                if self.token_rotator is None:
+                    raise BoltError(self._config_error_message)
+                refreshed = self.token_rotator.perform_bot_token_rotation(
+                    bot=bot,
+                    minutes_before_expiration=self.token_rotation_expiration_minutes,
+                )
+                if refreshed is not None:
+                    self.installation_store.save_bot(refreshed)
+                    bot_token = refreshed.bot_token
+
         token: Optional[str] = bot_token or user_token
         if token is None:
             # No valid token was found

diff --git a/slack_bolt/oauth/async_oauth_flow.py b/slack_bolt/oauth/async_oauth_flow.py
@@ -338,9 +338,13 @@ async def run_installation(self, code: str) -> Optional[Installation]:
                 bot_id=bot_id,
                 bot_user_id=oauth_response.get("bot_user_id"),
                 bot_scopes=oauth_response.get("scope"),  # comma-separated string
+                bot_refresh_token=oauth_response.get("refresh_token"),  # since v1.7
+                bot_token_expires_in=oauth_response.get("expires_in"),  # since v1.7
                 user_id=installer.get("id"),
                 user_token=installer.get("access_token"),
                 user_scopes=installer.get("scope"),  # comma-separated string
+                user_refresh_token=installer.get("refresh_token"),  # since v1.7
+                user_token_expires_in=installer.get("expires_in"),  # since v1.7
                 incoming_webhook_url=incoming_webhook.get("url"),
                 incoming_webhook_channel=incoming_webhook.get("channel"),
                 incoming_webhook_channel_id=incoming_webhook.get("channel_id"),

diff --git a/slack_bolt/oauth/async_oauth_settings.py b/slack_bolt/oauth/async_oauth_settings.py
@@ -41,6 +41,7 @@ class AsyncOAuthSettings:
     # Installation Management
     installation_store: AsyncInstallationStore
     installation_store_bot_only: bool
+    token_rotation_expiration_minutes: int
     authorize: AsyncAuthorize
     # state parameter related configurations
     state_store: AsyncOAuthStateStore
@@ -73,6 +74,7 @@ def __init__(
         # Installation Management
         installation_store: Optional[AsyncInstallationStore] = None,
         installation_store_bot_only: bool = False,
+        token_rotation_expiration_minutes: int = 120,
         # state parameter related configurations
         state_store: Optional[AsyncOAuthStateStore] = None,
         state_cookie_name: str = OAuthStateUtils.default_cookie_name,
@@ -140,8 +142,12 @@ def __init__(
             installation_store or get_or_create_default_installation_store(client_id)
         )
         self.installation_store_bot_only = installation_store_bot_only
+        self.token_rotation_expiration_minutes = token_rotation_expiration_minutes
         self.authorize = AsyncInstallationStoreAuthorize(
             logger=logger,
+            client_id=self.client_id,
+            client_secret=self.client_secret,
+            token_rotation_expiration_minutes=self.token_rotation_expiration_minutes,
             installation_store=self.installation_store,
             bot_only=self.installation_store_bot_only,
         )

diff --git a/slack_bolt/oauth/oauth_flow.py b/slack_bolt/oauth/oauth_flow.py
@@ -107,6 +107,7 @@ def sqlite3(
         state_cookie_name: str = OAuthStateUtils.default_cookie_name,
         state_expiration_seconds: int = OAuthStateUtils.default_expiration_seconds,
         installation_store_bot_only: bool = False,
+        token_rotation_expiration_minutes: int = 120,
         client: Optional[WebClient] = None,
         logger: Optional[Logger] = None,
     ) -> "OAuthFlow":
@@ -140,6 +141,7 @@ def sqlite3(
                     logger=logger,
                 ),
                 installation_store_bot_only=installation_store_bot_only,
+                token_rotation_expiration_minutes=token_rotation_expiration_minutes,
                 # state parameter related configurations
                 state_store=SQLite3OAuthStateStore(
                     database=database,
@@ -333,9 +335,13 @@ def run_installation(self, code: str) -> Optional[Installation]:
                 bot_id=bot_id,
                 bot_user_id=oauth_response.get("bot_user_id"),
                 bot_scopes=oauth_response.get("scope"),  # comma-separated string
+                bot_refresh_token=oauth_response.get("refresh_token"),  # since v1.7
+                bot_token_expires_in=oauth_response.get("expires_in"),  # since v1.7
                 user_id=installer.get("id"),
                 user_token=installer.get("access_token"),
                 user_scopes=installer.get("scope"),  # comma-separated string
+                user_refresh_token=installer.get("refresh_token"),  # since v1.7
+                user_token_expires_in=installer.get("expires_in"),  # since v1.7
                 incoming_webhook_url=incoming_webhook.get("url"),
                 incoming_webhook_channel=incoming_webhook.get("channel"),
                 incoming_webhook_channel_id=incoming_webhook.get("channel_id"),

diff --git a/slack_bolt/oauth/oauth_settings.py b/slack_bolt/oauth/oauth_settings.py
@@ -36,6 +36,7 @@ class OAuthSettings:
     # Installation Management
     installation_store: InstallationStore
     installation_store_bot_only: bool
+    token_rotation_expiration_minutes: int
     authorize: Authorize
     # state parameter related configurations
     state_store: OAuthStateStore
@@ -68,6 +69,7 @@ def __init__(
         # Installation Management
         installation_store: Optional[InstallationStore] = None,
         installation_store_bot_only: bool = False,
+        token_rotation_expiration_minutes: int = 120,
         # state parameter related configurations
         state_store: Optional[OAuthStateStore] = None,
         state_cookie_name: str = OAuthStateUtils.default_cookie_name,
@@ -134,8 +136,12 @@ def __init__(
             installation_store or get_or_create_default_installation_store(client_id)
         )
         self.installation_store_bot_only = installation_store_bot_only
+        self.token_rotation_expiration_minutes = token_rotation_expiration_minutes
         self.authorize = InstallationStoreAuthorize(
             logger=logger,
+            client_id=self.client_id,
+            client_secret=self.client_secret,
+            token_rotation_expiration_minutes=self.token_rotation_expiration_minutes,
             installation_store=self.installation_store,
             bot_only=self.installation_store_bot_only,
         )