diff --git a/aiogram/utils/web_app_signature.py b/aiogram/utils/web_app_signature.py index ffce0ead..623cc802 100644 --- a/aiogram/utils/web_app_signature.py +++ b/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") diff --git a/tests/test_utils/test_web_app_signature.py b/tests/test_utils/test_web_app_signature.py new file mode 100644 index 00000000..e8fbba71 --- /dev/null +++ b/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, + ) \ No newline at end of file