feat: added incoming requests verification using authorization JWT to…

…kens (#447)

README.md

@@ -103,7 +103,10 @@ app.add_event_handler("shutdown", bot.shutdown)
 # (сообщения и системные события).
 @app.post("/command")
 async def command_handler(request: Request) -> JSONResponse:
-    bot.async_execute_raw_bot_command(await request.json())
+    bot.async_execute_raw_bot_command(
+        await request.json(),
+        request_headers=request.headers,
+    )
     return JSONResponse(
         build_command_accepted_response(),
         status_code=HTTPStatus.ACCEPTED,
@@ -114,15 +117,21 @@ async def command_handler(request: Request) -> JSONResponse:
 # доступность бота и его список команд.
 @app.get("/status")
 async def status_handler(request: Request) -> JSONResponse:
-    status = await bot.raw_get_status(dict(request.query_params))
+    status = await bot.raw_get_status(
+        dict(request.query_params),
+        request_headers=request.headers,
+    )
     return JSONResponse(status)
 
 
 # На этот эндпоинт приходят коллбэки с результатами
 # выполнения асинхронных методов в BotX.
 @app.post("/notification/callback")
 async def callback_handler(request: Request) -> JSONResponse:
-    await bot.set_raw_botx_method_result(await request.json())
+    await bot.set_raw_botx_method_result(
+        await request.json(),
+        verify_request=False,
+    )
     return JSONResponse(
         build_command_accepted_response(),
         status_code=HTTPStatus.ACCEPTED,

pybotx/__init__.py

@@ -3,17 +3,25 @@
     UnsupportedBotAPIVersionError,
 )
 from pybotx.bot.api.responses.bot_disabled import (
+    BotAPIBotDisabledErrorData,
     BotAPIBotDisabledResponse,
     build_bot_disabled_response,
 )
 from pybotx.bot.api.responses.command_accepted import build_command_accepted_response
+from pybotx.bot.api.responses.unverified_request import (
+    BotAPIUnverifiedRequestErrorData,
+    BotAPIUnverifiedRequestResponse,
+    build_unverified_request_response,
+)
 from pybotx.bot.bot import Bot
 from pybotx.bot.callbacks.callback_repo_proto import CallbackRepoProto
 from pybotx.bot.exceptions import (
     AnswerDestinationLookupError,
     BotShuttingDownError,
     BotXMethodCallbackNotFoundError,
+    RequestHeadersNotProvidedError,
     UnknownBotAccountError,
+    UnverifiedRequestError,
 )
 from pybotx.bot.handler import IncomingMessageHandlerFunc, Middleware
 from pybotx.bot.handler_collector import HandlerCollector
@@ -121,14 +129,17 @@
 __all__ = (
     "AddedToChatEvent",
     "AnswerDestinationLookupError",
-    "AttachmentTypes",
     "AttachmentDocument",
     "AttachmentImage",
-    "AttachmentVoice",
+    "AttachmentTypes",
     "AttachmentVideo",
+    "AttachmentVoice",
     "Bot",
+    "BotAPIBotDisabledErrorData",
     "BotAPIBotDisabledResponse",
     "BotAPIMethodFailedCallback",
+    "BotAPIUnverifiedRequestErrorData",
+    "BotAPIUnverifiedRequestResponse",
     "BotAccount",
     "BotAccountWithSecret",
     "BotIsNotChatMemberError",
@@ -137,13 +148,13 @@
     "BotShuttingDownError",
     "BotXMethodCallbackNotFoundError",
     "BotXMethodFailedCallbackReceivedError",
+    "BotsListItem",
     "BubbleMarkup",
     "Button",
     "ButtonRow",
     "ButtonTextAlign",
     "CTSLoginEvent",
     "CTSLogoutEvent",
-    "EventEdit",
     "CallbackNotReceivedError",
     "CallbackRepoProto",
     "CantUpdatePersonalChatError",
@@ -161,6 +172,7 @@
     "DeletedFromChatEvent",
     "Document",
     "EditMessage",
+    "EventEdit",
     "EventNotFoundError",
     "File",
     "FileDeletedError",
@@ -198,11 +210,11 @@
     "RateLimitReachedError",
     "Reply",
     "ReplyMessage",
+    "RequestHeadersNotProvidedError",
     "SmartApp",
     "SmartAppEvent",
     "SmartAppEvent",
     "StatusRecipient",
-    "BotsListItem",
     "StealthModeDisabledError",
     "Sticker",
     "StickerPack",
@@ -211,6 +223,7 @@
     "UnknownBotAccountError",
     "UnknownSystemEventError",
     "UnsupportedBotAPIVersionError",
+    "UnverifiedRequestError",
     "UserDevice",
     "UserFromCSV",
     "UserFromSearch",
@@ -221,6 +234,7 @@
     "Voice",
     "build_bot_disabled_response",
     "build_command_accepted_response",
+    "build_unverified_request_response",
     "lifespan_wrapper",
 )
 

pybotx/bot/api/responses/bot_disabled.py

@@ -23,7 +23,7 @@ class BotAPIBotDisabledResponse:
 def build_bot_disabled_response(status_message: str) -> Dict[str, Any]:
     """Build bot disabled response for BotX.
 
-    It should be send if the bot can't process the command.
+    It should be sent if the bot can't process the command.
     If you would like to build complex response, see `BotAPIBotDisabledResponse`.
 
     :param status_message: Status message.

pybotx/bot/api/responses/unverified_request.py

@@ -0,0 +1,39 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, List, Literal
+
+
+@dataclass
+class BotAPIUnverifiedRequestErrorData:
+    status_message: str
+
+
+@dataclass
+class BotAPIUnverifiedRequestResponse:
+    """`Unverified request` response model.
+
+    Only `.error_data.status_message` attribute will be displayed to
+    user. Other attributes will be visible only in BotX logs.
+    """
+
+    error_data: BotAPIUnverifiedRequestErrorData
+    errors: List[str] = field(default_factory=list)
+    reason: Literal["unverified_request"] = "unverified_request"
+
+
+def build_unverified_request_response(status_message: str) -> Dict[str, Any]:
+    """Build `unverified request` response for BotX.
+
+    It should be sent if the header with the authorization token is missing or
+    the authorization token is invalid.
+    If you would like to build complex response, see `BotAPIUnverifiedRequestResponse`.
+
+    :param status_message: Status message.
+
+    :return: Built `unverified request` response.
+    """
+
+    response = BotAPIUnverifiedRequestResponse(
+        error_data=BotAPIUnverifiedRequestErrorData(status_message=status_message),
+    )
+
+    return asdict(response)

pybotx/bot/bot.py

@@ -9,6 +9,7 @@
     Dict,
     Iterator,
     List,
+    Mapping,
     Optional,
     Sequence,
     Tuple,
@@ -18,6 +19,7 @@
 
 import aiofiles
 import httpx
+import jwt
 from aiocsv.readers import AsyncDictReader
 from aiofiles.tempfile import NamedTemporaryFile, TemporaryDirectory
 from pydantic import ValidationError, parse_obj_as
@@ -28,7 +30,12 @@
 from pybotx.bot.callbacks.callback_memory_repo import CallbackMemoryRepo
 from pybotx.bot.callbacks.callback_repo_proto import CallbackRepoProto
 from pybotx.bot.contextvars import bot_id_var, chat_id_var
-from pybotx.bot.exceptions import AnswerDestinationLookupError
+from pybotx.bot.exceptions import (
+    AnswerDestinationLookupError,
+    RequestHeadersNotProvidedError,
+    UnknownBotAccountError,
+    UnverifiedRequestError,
+)
 from pybotx.bot.handler import Middleware
 from pybotx.bot.handler_collector import HandlerCollector
 from pybotx.bot.middlewares.exception_middleware import ExceptionHandlersDict
@@ -274,6 +281,8 @@ def __init__(
     def async_execute_raw_bot_command(
         self,
         raw_bot_command: Dict[str, Any],
+        verify_request: bool = True,
+        request_headers: Optional[Mapping[str, str]] = None,
         logging_command: bool = True,
     ) -> None:
         if logging_command:
@@ -284,6 +293,11 @@ def async_execute_raw_bot_command(
                 ),
             )
 
+        if verify_request:
+            if request_headers is None:
+                raise RequestHeadersNotProvidedError
+            self._verify_request(request_headers)
+
         try:
             bot_api_command: BotAPICommand = parse_obj_as(
                 # Same ignore as in pydantic
@@ -305,12 +319,22 @@ def async_execute_bot_command(
 
         return self._handler_collector.async_handle_bot_command(self, bot_command)
 
-    async def raw_get_status(self, query_params: Dict[str, str]) -> Dict[str, Any]:
+    async def raw_get_status(
+        self,
+        query_params: Dict[str, str],
+        verify_request: bool = True,
+        request_headers: Optional[Mapping[str, str]] = None,
+    ) -> Dict[str, Any]:
         logger.opt(lazy=True).debug(
             "Got status: {status}",
             status=lambda: pformat_jsonable_obj(query_params),
         )
 
+        if verify_request:
+            if request_headers is None:
+                raise RequestHeadersNotProvidedError
+            self._verify_request(request_headers)
+
         try:
             bot_api_status_recipient = BotAPIStatusRecipient.parse_obj(query_params)
         except ValidationError as exc:
@@ -330,9 +354,16 @@ async def get_status(self, status_recipient: StatusRecipient) -> BotMenu:
     async def set_raw_botx_method_result(
         self,
         raw_botx_method_result: Dict[str, Any],
+        verify_request: bool = True,
+        request_headers: Optional[Mapping[str, str]] = None,
     ) -> None:
         logger.debug("Got callback: {callback}", callback=raw_botx_method_result)
 
+        if verify_request:
+            if request_headers is None:
+                raise RequestHeadersNotProvidedError
+            self._verify_request(request_headers)
+
         callback: BotXMethodCallback = parse_obj_as(
             # Same ignore as in pydantic
             BotXMethodCallback,  # type: ignore[arg-type]
@@ -1909,6 +1940,48 @@ async def collect_metric(
         )
         await method.execute(payload)
 
+    def _verify_request(self, headers: Mapping[str, str]) -> None:  # noqa: WPS238
+        authorization_header = headers.get("authorization")
+        if not authorization_header:
+            raise UnverifiedRequestError("The authorization token was not provided.")
+
+        token = authorization_header.split()[-1]
+        decode_algorithms = ["HS256"]
+
+        try:
+            token_payload = jwt.decode(
+                jwt=token,
+                algorithms=decode_algorithms,
+                options={
+                    "verify_signature": False,
+                },
+            )
+        except jwt.DecodeError as decode_exc:
+            raise UnverifiedRequestError(decode_exc.args[0]) from decode_exc
+
+        audience = token_payload.get("aud")
+        if not audience or not isinstance(audience, Sequence) or len(audience) != 1:
+            raise UnverifiedRequestError("Invalid audience parameter was provided.")
+
+        try:
+            bot_account = self._bot_accounts_storage.get_bot_account(UUID(audience[-1]))
+        except UnknownBotAccountError as unknown_bot_exc:
+            raise UnverifiedRequestError(unknown_bot_exc.args[0]) from unknown_bot_exc
+
+        try:
+            jwt.decode(
+                jwt=token,
+                key=bot_account.secret_key,
+                algorithms=decode_algorithms,
+                issuer=bot_account.host,
+                leeway=0.5,
+                options={
+                    "verify_aud": False,
+                },
+            )
+        except jwt.InvalidTokenError as exc:
+            raise UnverifiedRequestError(exc.args[0]) from exc
+
     @staticmethod
     def _build_main_collector(
         collectors: Sequence[HandlerCollector],

pybotx/bot/bot_accounts_storage.py

@@ -13,11 +13,18 @@ def __init__(self, bot_accounts: List[BotAccountWithSecret]) -> None:
         self._bot_accounts = bot_accounts
         self._auth_tokens: Dict[UUID, str] = {}
 
+    def get_bot_account(self, bot_id: UUID) -> BotAccountWithSecret:
+        for bot_account in self._bot_accounts:
+            if bot_account.id == bot_id:
+                return bot_account
+
+        raise UnknownBotAccountError(bot_id)
+
     def iter_bot_accounts(self) -> Iterator[BotAccount]:
         yield from self._bot_accounts
 
     def get_host(self, bot_id: UUID) -> str:
-        bot_account = self._get_bot_account(bot_id)
+        bot_account = self.get_bot_account(bot_id)
         return bot_account.host
 
     def set_token(self, bot_id: UUID, token: str) -> None:
@@ -27,7 +34,7 @@ def get_token_or_none(self, bot_id: UUID) -> Optional[str]:
         return self._auth_tokens.get(bot_id)
 
     def build_signature(self, bot_id: UUID) -> str:
-        bot_account = self._get_bot_account(bot_id)
+        bot_account = self.get_bot_account(bot_id)
 
         signed_bot_id = hmac.new(
             key=bot_account.secret_key.encode(),
@@ -38,11 +45,4 @@ def build_signature(self, bot_id: UUID) -> str:
         return base64.b16encode(signed_bot_id).decode()
 
     def ensure_bot_id_exists(self, bot_id: UUID) -> None:
-        self._get_bot_account(bot_id)
-
-    def _get_bot_account(self, bot_id: UUID) -> BotAccountWithSecret:
-        for bot_account in self._bot_accounts:
-            if bot_account.id == bot_id:
-                return bot_account
-
-        raise UnknownBotAccountError(bot_id)
+        self.get_bot_account(bot_id)

pybotx/bot/exceptions.py

@@ -29,3 +29,14 @@ class AnswerDestinationLookupError(Exception):
     def __init__(self) -> None:
         self.message = "No IncomingMessage received. Use `Bot.send` instead"
         super().__init__(self.message)
+
+
+class RequestHeadersNotProvidedError(Exception):
+    def __init__(self, *args: Any) -> None:
+        reason = "To verify the request you should provide headers."
+        message = args[0] if args else reason
+        super().__init__(message)
+
+
+class UnverifiedRequestError(Exception):
+    """The authorization header is missing or the token is invalid."""