Repository: swimmwatch/telegram-webapp-auth

Issue:
Add `validate_third_party` method
Desciption: https://core.telegram.org/bots/webapps#validating-data-for-third-party-use

Also implement Init Data object.

Discussion:

Solution patch:
diff --git a/telegram_webapp_auth/auth.py b/telegram_webapp_auth/auth.py
index beb8447..0a92d33 100644
--- a/telegram_webapp_auth/auth.py
+++ b/telegram_webapp_auth/auth.py
@@ -1,3 +1,4 @@
+import base64
 import hashlib
 import hmac
 import json
@@ -11,6 +12,8 @@
 from cryptography.exceptions import InvalidSignature
 from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
 
+from telegram_webapp_auth.data import PROD_PUBLIC_KEY
+from telegram_webapp_auth.data import TEST_PUBLIC_KEY
 from telegram_webapp_auth.data import WebAppChat
 from telegram_webapp_auth.data import WebAppInitData
 from telegram_webapp_auth.data import WebAppUser
@@ -40,7 +43,7 @@ def __init__(self, secret: bytes) -> None:
         self._secret = secret
 
     @staticmethod
-    def _parse_init_data(data: str) -> dict:
+    def __parse_init_data(data: str) -> dict:
         """Convert init_data string into dictionary.
 
         Args:
@@ -52,7 +55,7 @@ def _parse_init_data(data: str) -> dict:
         return dict(param.split("=") for param in data.split("&"))
 
     @staticmethod
-    def _parse_json(data: str) -> dict:
+    def __parse_json(data: str) -> dict:
         """Convert JSON string value from WebAppInitData to Python dictionary.
 
         Links:
@@ -66,7 +69,7 @@ def _parse_json(data: str) -> dict:
         except JSONDecodeError:
             raise InvalidInitDataError("Cannot decode init data")
 
-    def _validate(self, hash_: str, token: str) -> bool:
+    def _validate(self, hash_: str, init_data: str) -> bool:
         """Validates the data received from the Telegram web app, using the method from Telegram documentation.
 
         Links:
@@ -74,22 +77,22 @@ def _validate(self, hash_: str, token: str) -> bool:
 
         Args:
             hash_: hash from init data
-            token: init data from webapp
+            init_data: init data from webapp
 
         Returns:
             bool: Validation result
         """
-        token_bytes = token.encode("utf-8")
-        client_hash = hmac.new(self._secret, token_bytes, hashlib.sha256).hexdigest()
+        init_data_bytes = init_data.encode("utf-8")
+        client_hash = hmac.new(self._secret, init_data_bytes, hashlib.sha256).hexdigest()
         return hmac.compare_digest(client_hash, hash_)
 
     @staticmethod
-    def _ed25519_verify(
+    def __ed25519_verify(
         public_key: Ed25519PublicKey,
         signature: bytes,
         message: bytes,
     ) -> bool:
-        """Verify the signature of the message using the public key.
+        """Verify the Ed25519 signature of the message using the public key.
 
         Args:
             public_key: public key
@@ -106,6 +109,119 @@ def _ed25519_verify(
         except InvalidSignature:
             return False
 
+    @staticmethod
+    def __decode_signature(val: str):
+        """Decode a base64-encoded signature, appending padding if necessary.
+
+        :param val: A base64-encoded string.
+        :return: Decoded signature as bytes.
+        """
+
+        # Add padding if the length is not a multiple of 4
+        padded_v = val + "=" * ((4 - len(val) % 4) % 4)
+
+        # Decode the Base64 string
+        try:
+            signature = base64.urlsafe_b64decode(padded_v)
+            return signature
+        except Exception as err:
+            raise InvalidInitDataError(f"An error occurred during decoding: {err}")
+
+    def __serialize_init_data(self, init_data_dict: typing.Dict[str, typing.Any]) -> WebAppInitData:
+        """Serialize the init data dictionary into WebAppInitData object.
+
+        Args:
+            init_data_dict: the init data dictionary
+
+        Returns:
+            WebAppInitData: the serialized WebAppInitData object
+        """
+
+        user_data = init_data_dict.get("user")
+        if user_data:
+            user_data = self.__parse_json(user_data)
+
+        chat_data = init_data_dict.get("chat")
+        if chat_data:
+            chat_data = self.__parse_json(chat_data)
+
+        receiver_data = init_data_dict.get("receiver")
+        if receiver_data:
+            receiver_data = self.__parse_json(receiver_data)
+
+        data = init_data_dict | {
+            "user": WebAppUser(**user_data) if user_data else None,
+            "receiver": WebAppUser(**receiver_data) if receiver_data else None,
+            "chat": WebAppChat(**chat_data) if chat_data else None,
+        }
+        return WebAppInitData(**data)
+
+    def validate_third_party(
+        self,
+        init_data: str,
+        bot_id: int,
+        expr_in: typing.Optional[timedelta] = None,
+        is_test: bool = False,
+    ) -> WebAppInitData:
+        """Validates the data for Third-Party Use, using the method from Telegram documentation.
+
+        Links:
+            https://core.telegram.org/bots/webapps#validating-data-for-third-party-use
+
+        Args:
+            init_data: init data from mini app
+            bot_id: Telegram Bot ID
+            expr_in: time delta to check if the token is expired
+            is_test: true if the init data was issued in Telegram production environment
+
+        Returns:
+            WebAppInitData: parsed init a data object
+
+        Raises:
+            InvalidInitDataError: if the init data is invalid
+            ExpiredInitDataError: if the init data is expired
+        """
+        init_data = unquote(init_data)
+        init_data_dict = self.__parse_init_data(init_data)
+        data_check_string = "\n".join(
+            f"{key}={val}"
+            for key, val in sorted(init_data_dict.items(), key=lambda item: item[0])
+            if key != "hash" and key != "signature"
+        )
+        data_check_string = f"{bot_id}:WebAppData\n{data_check_string}"
+
+        signature = init_data_dict.get("signature")
+        if not signature:
+            raise InvalidInitDataError("Init data does not contain signature")
+
+        signature = signature.strip()
+
+        if is_test:
+            public_key = TEST_PUBLIC_KEY
+        else:
+            public_key = PROD_PUBLIC_KEY
+
+        signature_bytes = self.__decode_signature(signature)
+        data_check_string_bytes = data_check_string.encode("utf-8")
+        if not self.__ed25519_verify(public_key, signature_bytes, data_check_string_bytes):
+            raise InvalidInitDataError("Invalid data")
+
+        auth_date = init_data_dict.get("auth_date")
+        if not auth_date:
+            raise InvalidInitDataError("Init data does not contain auth_date")
+
+        try:
+            auth_dt = datetime.fromtimestamp(float(auth_date), tz=timezone.utc)
+        except ValueError:
+            raise InvalidInitDataError("Invalid auth_date")
+
+        if expr_in:
+            now = datetime.now(tz=timezone.utc)
+            if now - auth_dt > expr_in:
+                raise ExpiredInitDataError
+
+        return self.__serialize_init_data(init_data_dict)
+
     def validate(
         self,
         init_data: str,
@@ -128,7 +244,7 @@ def validate(
             ExpiredInitDataError: if the init data is expired
         """
         init_data = unquote(init_data)
-        init_data_dict = self._parse_init_data(init_data)
+        init_data_dict = self.__parse_init_data(init_data)
         data_check_string = "\n".join(
             f"{key}={val}" for key, val in sorted(init_data_dict.items(), key=lambda item: item[0]) if key != "hash"
         )
@@ -139,7 +255,7 @@ def validate(
         hash_ = hash_.strip()
 
         if not self._validate(hash_, data_check_string):
-            raise InvalidInitDataError("Invalid token")
+            raise InvalidInitDataError("Invalid data")
 
         auth_date = init_data_dict.get("auth_date")
         if not auth_date:
@@ -155,21 +271,4 @@ def validate(
             if now - auth_dt > expr_in:
                 raise ExpiredInitDataError
 
-        user_data = init_data_dict.get("user")
-        if user_data:
-            user_data = self._parse_json(user_data)
-
-        chat_data = init_data_dict.get("chat")
-        if chat_data:
-            chat_data = self._parse_json(chat_data)
-
-        receiver_data = init_data_dict.get("receiver")
-        if receiver_data:
-            receiver_data = self._parse_json(receiver_data)
-
-        data = init_data_dict | {
-            "user": WebAppUser(**user_data) if user_data else None,
-            "receiver": WebAppUser(**receiver_data) if receiver_data else None,
-            "chat": WebAppChat(**chat_data) if chat_data else None,
-        }
-        return WebAppInitData(**data)
+        return self.__serialize_init_data(init_data_dict)


Test patch:
diff --git a/tests/test_validate_third_party.py b/tests/test_validate_third_party.py
new file mode 100644
index 0000000..f3120e6
--- /dev/null
+++ b/tests/test_validate_third_party.py
@@ -0,0 +1,79 @@
+import typing
+from datetime import timedelta
+
+import pytest
+
+from telegram_webapp_auth.auth import TelegramAuthenticator
+from telegram_webapp_auth.errors import ExpiredInitDataError
+from telegram_webapp_auth.errors import InvalidInitDataError
+
+_TEST_INIT_DATA = (
+    "query_id=AAH1t3EVAAAAAPW3cRVyuBgH&"
+    "user=%7B%22id%22%3A359774197%2C%22first_name%22%3A%22Dmitry%22%2C%22"
+    "last_name%22%3A%22Vasiliev%22%2C%22language_code%22%3A%22en%22%2C%22"
+    "allows_write_to_pm%22%3Atrue%2C%22"
+    "photo_url%22%3A%22https%3A%5C%2F%5C%2Ft.me%5C%2Fi%5C%2Fuserpic%5C%2F320%5C%2FaUkVEo2bVNr6G41sIB2bNwCdbkwnaVS79N41WHr_ApQ.svg%22%7D&"  # noqa: E501
+    "auth_date=1736353840&"
+    "signature=s72bv8J1hwJanbDqlo9TTMK6Uf4WSwQKuPKK_Q16QBhKD0hfOfoYCOpRl_d8m_8AEI1_oF-9WCJuwW1KQy5-BA&"
+    "hash=c0b008672bf1edae403293e790eec42d047ac6613a5799e1939fac5eca41a397"
+)
+_TEST_BOT_ID = 7544535829
+
+
+@pytest.mark.parametrize(
+    "test_input,expected_err",
+    [
+        # Test case 1: invalid format (semicolon added)
+        (
+            ";" + _TEST_INIT_DATA,
+            InvalidInitDataError,
+        ),
+        # Test case 2: positive case
+        (
+            _TEST_INIT_DATA,
+            None,
+        ),
+    ],
+)
+def test_parse(
+    test_input: str,
+    expected_err: typing.Optional[typing.Type[Exception]],
+    authenticator: TelegramAuthenticator,
+) -> None:
+    if expected_err:
+        with pytest.raises(expected_err):
+            authenticator.validate_third_party(test_input, _TEST_BOT_ID)
+    else:
+        authenticator.validate_third_party(test_input, _TEST_BOT_ID)
+
+
+@pytest.mark.parametrize(
+    "expr_in,expected_err",
+    [
+        # Test case 1: valid input data
+        (
+            None,
+            None,
+        ),
+        # Test case 2: expired on 1 second
+        (
+            timedelta(seconds=1),
+            ExpiredInitDataError,
+        ),
+        # Test case 3: not expired
+        (
+            timedelta(seconds=0),
+            None,
+        ),
+    ],
+)
+def test_parse_expire(
+    authenticator: TelegramAuthenticator,
+    expr_in: typing.Optional[timedelta],
+    expected_err: typing.Optional[typing.Type[Exception]],
+) -> None:
+    if expected_err:
+        with pytest.raises(expected_err):
+            authenticator.validate_third_party(_TEST_INIT_DATA, _TEST_BOT_ID, expr_in)
+    else:
+        authenticator.validate_third_party(_TEST_INIT_DATA, _TEST_BOT_ID, expr_in)