Skip to content

Commit

Permalink
Another approach to bot id
Browse files Browse the repository at this point in the history
  • Loading branch information
coder2020official committed Jul 26, 2024
1 parent 536ffa2 commit d2485bf
Show file tree
Hide file tree
Showing 4 changed files with 42 additions and 20 deletions.
22 changes: 13 additions & 9 deletions telebot/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,9 +152,14 @@ class TeleBot:
:param allow_sending_without_reply: Default value for allow_sending_without_reply, defaults to None
:type allow_sending_without_reply: :obj:`bool`, optional
:param colorful_logs: Outputs colorful logs
:type colorful_logs: :obj:`bool`, optional
:param validate_token: Validate token, defaults to True;
:type validate_token: :obj:`bool`, optional
:raises ImportError: If coloredlogs module is not installed and colorful_logs is True
:raises ValueError: If token is invalid
"""

def __init__(
Expand All @@ -169,7 +174,7 @@ def __init__(
protect_content: Optional[bool]=None,
allow_sending_without_reply: Optional[bool]=None,
colorful_logs: Optional[bool]=False,
token_check: Optional[bool]=True
validate_token: Optional[bool]=True
):

# update-related
Expand All @@ -186,11 +191,12 @@ def __init__(
self.allow_sending_without_reply = allow_sending_without_reply
self.webhook_listener = None
self._user = None
self.bot_id: int = None

# token check
if token_check:
self._user = self.get_me()
self.bot_id = self._user.id
if validate_token:
util.validate_token(self.token)

self.bot_id = util.extract_bot_id(self.token) # subject to change in future, unspecified

# logs-related
if colorful_logs:
Expand Down Expand Up @@ -286,9 +292,7 @@ def __init__(
self.threaded = threaded
if self.threaded:
self.worker_pool = util.ThreadPool(self, num_threads=num_threads)




@property
def user(self) -> types.User:
"""
Expand Down
20 changes: 11 additions & 9 deletions telebot/async_telebot.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,8 +117,11 @@ class AsyncTeleBot:
:param colorful_logs: Outputs colorful logs
:type colorful_logs: :obj:`bool`, optional
:param token_check: Check token on start
:type token_check: :obj:`bool`, optional, defaults to True
:param validate_token: Validate token, defaults to True;
:type validate_token: :obj:`bool`, optional
:raises ImportError: If coloredlogs module is not installed and colorful_logs is True
:raises ValueError: If token is invalid
"""

def __init__(self, token: str, parse_mode: Optional[str]=None, offset: Optional[int]=None,
Expand All @@ -129,7 +132,7 @@ def __init__(self, token: str, parse_mode: Optional[str]=None, offset: Optional[
protect_content: Optional[bool]=None,
allow_sending_without_reply: Optional[bool]=None,
colorful_logs: Optional[bool]=False,
token_check: Optional[bool]=True) -> None:
validate_token: Optional[bool]=True) -> None:

# update-related
self.token = token
Expand Down Expand Up @@ -186,15 +189,14 @@ def __init__(self, token: str, parse_mode: Optional[str]=None, offset: Optional[
self.middlewares = []

self._user = None # set during polling
self.bot_id = None
self.bot_id: int = None

if token_check:
result = apihelper.get_me(token)
self._user = types.User.de_json(result)
self.bot_id = self._user.id
if validate_token:
util.validate_token(self.token)

self.bot_id: int = util.extract_bot_id(self.token) # subject to change, unspecified



@property
def user(self):
return self._user
Expand Down
18 changes: 17 additions & 1 deletion telebot/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -686,6 +686,22 @@ def validate_web_app_data(token: str, raw_init_data: str):
return hmac.new(secret_key.digest(), data_check_string.encode(), sha256).hexdigest() == init_data_hash


def validate_token(token) -> bool:
if any(char.isspace() for char in token):
raise ValueError('Token must not contain spaces')

if ':' not in token:
raise ValueError('Token must contain a colon')

if len(token.split(':')) != 2:
raise ValueError('Token must contain exactly 2 parts separated by a colon')

return True

def extract_bot_id(token) -> str:
return int(token.split(':')[0])


__all__ = (
"content_type_media", "content_type_service", "update_types",
"WorkerThread", "AsyncTask", "CustomRequestResponse",
Expand All @@ -696,5 +712,5 @@ def validate_web_app_data(token: str, raw_init_data: str):
"split_string", "smart_split", "escape", "user_link", "quick_markup",
"antiflood", "parse_web_app_data", "validate_web_app_data",
"or_set", "or_clear", "orify", "OrEvent", "per_thread",
"webhook_google_functions"
"webhook_google_functions", "validate_token", "extract_bot_id"
)
2 changes: 1 addition & 1 deletion tests/test_handler_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

@pytest.fixture()
def telegram_bot():
return telebot.TeleBot('', threaded=False, token_check=False)
return telebot.TeleBot('', threaded=False, validate_token=False)


@pytest.fixture
Expand Down

0 comments on commit d2485bf

Please sign in to comment.