Add tests.

aiogram/utils/web_app_signature.py

@@ -13,7 +13,9 @@ PRODUCTION_PUBLIC_KEY = bytes.fromhex(
 TEST_PUBLIC_KEY = bytes.fromhex("40055058a4ee38156a06562e52eece92a771bcd8346a8c4615cb7376eddf72ec")
 
 
-def check_webapp_signature(bot_id: int, init_data: str, is_test: bool = False) -> bool:
+def check_webapp_signature(
+    bot_id: int, init_data: str, public_key_bytes: bytes = PRODUCTION_PUBLIC_KEY
+) -> bool:
     """
     Check incoming WebApp init data signature without bot token using only bot id.
 
@@ -21,7 +23,7 @@ def check_webapp_signature(bot_id: int, init_data: str, is_test: bool = False) -
 
     :param bot_id: Bot ID
     :param init_data: WebApp init data
-    :param is_test: Is test environment
+    :param public_key: Public key
     :return: True if signature is valid, False otherwise
     """
     try:
@@ -43,7 +45,6 @@ def check_webapp_signature(bot_id: int, init_data: str, is_test: bool = False) -
     padding = "=" * (-len(signature_b64) % 4)
     signature = base64.urlsafe_b64decode(signature_b64 + padding)
 
-    public_key_bytes = TEST_PUBLIC_KEY if is_test else PRODUCTION_PUBLIC_KEY
     public_key = Ed25519PublicKey.from_public_bytes(public_key_bytes)
 
     try:
@@ -54,16 +55,16 @@ def check_webapp_signature(bot_id: int, init_data: str, is_test: bool = False) -
 
 
 def safe_check_webapp_init_data_from_signature(
-    bot_id: int, init_data: str, is_test: bool = False
+    bot_id: int, init_data: str, public_key_bytes: bytes = PRODUCTION_PUBLIC_KEY
 ) -> WebAppInitData:
     """
     Validate raw WebApp init data using only bot id and return it as WebAppInitData object
 
     :param bot_id: bot id
     :param init_data: data from frontend to be parsed and validated
-    :param is_test: is test environment, default is False
+    :param public_key_bytes: public key
     :return: WebAppInitData object
     """
-    if check_webapp_signature(bot_id, init_data, is_test):
+    if check_webapp_signature(bot_id, init_data, public_key_bytes):
         return parse_webapp_init_data(init_data)
     raise ValueError("Invalid init data signature")

tests/test_utils/test_web_app_signature.py

@@ -0,0 +1,63 @@
+import pytest
+
+from aiogram.utils.web_app import WebAppInitData
+from aiogram.utils.web_app_signature import (
+    check_webapp_signature,
+    safe_check_webapp_init_data_from_signature,
+)
+
+PRIVATE_KEY = bytes.fromhex("c80e09dc60f5efcf2e1f8d0793358e0ea3371267bef0024588f7bf67cf48dfb9")
+PUBLIC_KEY = bytes.fromhex("4112765021341e5415e772cd65903f6b94e3ea1c2ab669e6d3e18ee2db00da61")
+
+
+class TestWebAppSignature:
+    @pytest.mark.parametrize(
+        "bot_id,case,result",
+        [
+            [
+                42,
+                "auth_date=1650385342&user=%7B%22id%22%3A42%2C%22first_name%22%3A%22Test%22%7D&query_id=test&signature=JQ0JR2tjC65yq_jNZV0wuJVX6J-SWPMV0mprUXG34g-NvxL4RcF1Rz5n4VVo00VRghEUBf5t___uoeb1-jU_Cw",
+                True,
+            ],
+            [
+                42,
+                "auth_date=1650385342&user=%7B%22id%22%3A42%2C%22first_name%22%3A%22Test%22%7D&query_id=test&signature=JQ0JR2tjC65yq_jNZV0wuJVX6J-SWPMV0mprUXG34g-NvxL4RcF1Rz5n4VVo00VRghEUBf5t___uoeb1-j1U_w",
+                False,
+            ],
+            [
+                42,
+                "auth_date=1650385342&user=%7B%22id%22%3A42%2C%22first_name%22%3A%22Test%22%7D&query_id=test",
+                False,
+            ],
+            [
+                42,
+                "",
+                False,
+            ],
+            [42, "test&foo=bar=baz", False],
+        ],
+    )
+    def test_check_webapp_signature(self, bot_id: int, case: str, result: bool):
+        assert check_webapp_signature(bot_id, case, PUBLIC_KEY) is result
+
+    def test_safe_check_webapp_init_data_from_signature(self):
+        result = safe_check_webapp_init_data_from_signature(
+            42,
+            "auth_date=1650385342&user=%7B%22id%22%3A42%2C%22first_name%22%3A%22Test%22%7D&query_id=test&hash=123&signature=JQ0JR2tjC65yq_jNZV0wuJVX6J-SWPMV0mprUXG34g-NvxL4RcF1Rz5n4VVo00VRghEUBf5t___uoeb1-jU_Cw",
+            PUBLIC_KEY,
+        )
+        assert isinstance(result, WebAppInitData)
+        assert result.user is not None
+        assert result.user.id == 42
+        assert result.user.first_name == "Test"
+        assert result.query_id == "test"
+        assert result.auth_date.year == 2022
+        assert result.hash == "123"
+
+    def test_safe_check_webapp_init_data_from_signature_invalid(self):
+        with pytest.raises(ValueError):
+            safe_check_webapp_init_data_from_signature(
+                42,
+                "auth_date=1650385342&user=%7B%22id%22%3A42%2C%22first_name%22%3A%22Test%22%7D&query_id=test&hash=123&signature=JQ0JR2tjC65yq_jNZV0wuJVX6J-SWPMV0mprUXG34g-NvxL4RcF1Rz5n4VVo00VRghEUBf5t___uoeb1-j1U_w",
+                PUBLIC_KEY,
+            )