From 08f0635afe07e0b9855f05526924af93ba45dd20 Mon Sep 17 00:00:00 2001 From: Oleg A Date: Wed, 28 Apr 2021 01:24:34 +0300 Subject: [PATCH] Deep linking util fix (#569) * fix: deep linking util fixed and refactored * fix: wrong payload split * feat: check payload length --- aiogram/utils/deep_linking.py | 131 +++++++++++++++++--------- tests/test_utils/test_deep_linking.py | 37 +++++--- 2 files changed, 110 insertions(+), 58 deletions(-) diff --git a/aiogram/utils/deep_linking.py b/aiogram/utils/deep_linking.py index acb105da..e8035d9a 100644 --- a/aiogram/utils/deep_linking.py +++ b/aiogram/utils/deep_linking.py @@ -1,10 +1,10 @@ """ Deep linking -Telegram bots have a deep linking mechanism, that allows for passing additional -parameters to the bot on startup. It could be a command that launches the bot — or -an auth token to connect the user's Telegram account to their account on some -external service. +Telegram bots have a deep linking mechanism, that allows for passing +additional parameters to the bot on startup. It could be a command that +launches the bot — or an auth token to connect the user's Telegram +account to their account on some external service. You can read detailed description in the source: https://core.telegram.org/bots#deep-linking @@ -16,86 +16,123 @@ Basic link example: .. code-block:: python from aiogram.utils.deep_linking import get_start_link - link = await get_start_link('foo') # result: 'https://t.me/MyBot?start=foo' + link = await get_start_link('foo') + + # result: 'https://t.me/MyBot?start=foo' Encoded link example: .. code-block:: python - from aiogram.utils.deep_linking import get_start_link, decode_payload - link = await get_start_link('foo', encode=True) # result: 'https://t.me/MyBot?start=Zm9v' - # and decode it back: - payload = decode_payload('Zm9v') # result: 'foo' + from aiogram.utils.deep_linking import get_start_link + + link = await get_start_link('foo', encode=True) + # result: 'https://t.me/MyBot?start=Zm9v' + +Decode it back example: + .. code-block:: python + + from aiogram.utils.deep_linking import decode_payload + from aiogram.types import Message + + @dp.message_handler(commands=["start"]) + async def handler(message: Message): + args = message.get_args() + payload = decode_payload(args) + await message.answer(f"Your payload: {payload}") """ +import re +from base64 import urlsafe_b64decode, urlsafe_b64encode + +from ..bot import Bot + +BAD_PATTERN = re.compile(r"[^_A-z0-9-]") async def get_start_link(payload: str, encode=False) -> str: """ - Use this method to handy get 'start' deep link with your payload. - If you need to encode payload or pass special characters - set encode as True + Get 'start' deep link with your payload. + + If you need to encode payload or pass special characters - + set encode as True :param payload: args passed with /start :param encode: encode payload with base64url :return: link """ - return await _create_link('start', payload, encode) + return await _create_link( + link_type="start", + payload=payload, + encode=encode, + ) async def get_startgroup_link(payload: str, encode=False) -> str: """ - Use this method to handy get 'startgroup' deep link with your payload. - If you need to encode payload or pass special characters - set encode as True + Get 'startgroup' deep link with your payload. + + If you need to encode payload or pass special characters - + set encode as True :param payload: args passed with /start :param encode: encode payload with base64url :return: link """ - return await _create_link('startgroup', payload, encode) + return await _create_link( + link_type="startgroup", + payload=payload, + encode=encode, + ) async def _create_link(link_type, payload: str, encode=False): + """ + Create deep link. + + :param link_type: `start` or `startgroup` + :param payload: any string-convertible data + :param encode: pass True to encode the payload + :return: deeplink + """ bot = await _get_bot_user() - payload = filter_payload(payload) - if encode: - payload = encode_payload(payload) - return f'https://t.me/{bot.username}?{link_type}={payload}' - -def encode_payload(payload: str) -> str: - """ Encode payload with URL-safe base64url. """ - from base64 import urlsafe_b64encode - result: bytes = urlsafe_b64encode(payload.encode()) - return result.decode() - - -def decode_payload(payload: str) -> str: - """ Decode payload with URL-safe base64url. """ - from base64 import urlsafe_b64decode - result: bytes = urlsafe_b64decode(payload + '=' * (4 - len(payload) % 4)) - return result.decode() - - -def filter_payload(payload: str) -> str: - """ Convert payload to text and search for not allowed symbols. """ - import re - - # convert to string if not isinstance(payload, str): payload = str(payload) - # search for not allowed characters - if re.search(r'[^_A-z0-9-]', payload): - message = ('Wrong payload! Only A-Z, a-z, 0-9, _ and - are allowed. ' - 'We recommend to encode parameters with binary and other ' - 'types of content.') + if encode: + payload = encode_payload(payload) + + if re.search(BAD_PATTERN, payload): + message = ( + "Wrong payload! Only A-Z, a-z, 0-9, _ and - are allowed. " + "Pass `encode=True` or encode payload manually." + ) raise ValueError(message) - return payload + if len(payload) > 64: + message = "Payload must be up to 64 characters long." + raise ValueError(message) + + return f"https://t.me/{bot.username}?{link_type}={payload}" + + +def encode_payload(payload: str) -> str: + """Encode payload with URL-safe base64url.""" + payload = str(payload) + bytes_payload: bytes = urlsafe_b64encode(payload.encode()) + str_payload = bytes_payload.decode() + return str_payload.replace("=", "") + + +def decode_payload(payload: str) -> str: + """Decode payload with URL-safe base64url.""" + payload += "=" * (4 - len(payload) % 4) + result: bytes = urlsafe_b64decode(payload) + return result.decode() async def _get_bot_user(): - """ Get current user of bot. """ - from ..bot import Bot + """Get current user of bot.""" bot = Bot.get_current() return await bot.me diff --git a/tests/test_utils/test_deep_linking.py b/tests/test_utils/test_deep_linking.py index a1d01e4e..f4aa14f1 100644 --- a/tests/test_utils/test_deep_linking.py +++ b/tests/test_utils/test_deep_linking.py @@ -1,7 +1,11 @@ import pytest -from aiogram.utils.deep_linking import decode_payload, encode_payload, filter_payload -from aiogram.utils.deep_linking import get_start_link, get_startgroup_link +from aiogram.utils.deep_linking import ( + decode_payload, + encode_payload, + get_start_link, + get_startgroup_link, +) from tests.types import dataset # enable asyncio mode @@ -17,9 +21,11 @@ PAYLOADS = [ WRONG_PAYLOADS = [ '@BotFather', + "Some:special$characters#=", 'spaces spaces spaces', 1234567890123456789.0, ] +USERNAME = dataset.USER["username"] @pytest.fixture(params=PAYLOADS, name='payload') @@ -47,7 +53,7 @@ def get_bot_user_fixture(monkeypatch): class TestDeepLinking: async def test_get_start_link(self, payload): link = await get_start_link(payload) - assert link == f'https://t.me/{dataset.USER["username"]}?start={payload}' + assert link == f'https://t.me/{USERNAME}?start={payload}' async def test_wrong_symbols(self, wrong_payload): with pytest.raises(ValueError): @@ -55,20 +61,29 @@ class TestDeepLinking: async def test_get_startgroup_link(self, payload): link = await get_startgroup_link(payload) - assert link == f'https://t.me/{dataset.USER["username"]}?startgroup={payload}' + assert link == f'https://t.me/{USERNAME}?startgroup={payload}' async def test_filter_encode_and_decode(self, payload): - _payload = filter_payload(payload) - encoded = encode_payload(_payload) + encoded = encode_payload(payload) decoded = decode_payload(encoded) assert decoded == str(payload) - async def test_get_start_link_with_encoding(self, payload): + async def test_get_start_link_with_encoding(self, wrong_payload): # define link - link = await get_start_link(payload, encode=True) + link = await get_start_link(wrong_payload, encode=True) # define reference link - payload = filter_payload(payload) - encoded_payload = encode_payload(payload) + encoded_payload = encode_payload(wrong_payload) - assert link == f'https://t.me/{dataset.USER["username"]}?start={encoded_payload}' + assert link == f'https://t.me/{USERNAME}?start={encoded_payload}' + + async def test_64_len_payload(self): + payload = "p" * 64 + link = await get_start_link(payload) + assert link + + async def test_too_long_payload(self): + payload = "p" * 65 + print(payload, len(payload)) + with pytest.raises(ValueError): + await get_start_link(payload)