From ad0f8284f686b2501460a51943b2149deb9331ce Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Thu, 1 Aug 2019 19:07:08 +0200 Subject: [PATCH] Use the correct way to parse peer identifiers --- pyrogram/client/client.py | 45 +++++++------- pyrogram/client/ext/utils.py | 60 ++++++++++++------- pyrogram/client/methods/chats/get_chat.py | 4 +- pyrogram/client/methods/chats/get_dialogs.py | 6 +- .../bots_and_keyboards/callback_query.py | 12 +--- .../types/messages_and_media/message.py | 7 ++- pyrogram/client/types/user_and_chats/chat.py | 5 +- .../client/types/user_and_chats/dialog.py | 12 +--- 8 files changed, 80 insertions(+), 71 deletions(-) diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 64595a91..511dbf2a 100644 --- a/pyrogram/client/client.py +++ b/pyrogram/client/client.py @@ -1015,7 +1015,7 @@ class Client(Methods, BaseClient): access_hash = 0 peer_type = "group" elif isinstance(peer, (types.Channel, types.ChannelForbidden)): - peer_id = int("-100" + str(peer.id)) + peer_id = utils.get_channel_id(peer.id) access_hash = peer.access_hash username = getattr(peer, "username", None) @@ -1131,7 +1131,7 @@ class Client(Methods, BaseClient): try: diff = self.send( functions.updates.GetChannelDifference( - channel=self.resolve_peer(int("-100" + str(channel_id))), + channel=self.resolve_peer(utils.get_channel_id(channel_id)), filter=types.ChannelMessagesFilter( ranges=[types.MessageRange( min_id=update.message.id, @@ -1519,33 +1519,38 @@ class Client(Methods, BaseClient): except KeyError: raise PeerIdInvalid - if peer_id > 0: + peer_type = utils.get_type(peer_id) + + if peer_type == "user": self.fetch_peers( self.send( functions.users.GetUsers( - id=[types.InputUser( - user_id=peer_id, - access_hash=0 - )] + id=[ + types.InputUser( + user_id=peer_id, + access_hash=0 + ) + ] ) ) ) + elif peer_type == "chat": + self.send( + functions.messages.GetChats( + id=[-peer_id] + ) + ) else: - if str(peer_id).startswith("-100"): - self.send( - functions.channels.GetChannels( - id=[types.InputChannel( - channel_id=int(str(peer_id)[4:]), + self.send( + functions.channels.GetChannels( + id=[ + types.InputChannel( + channel_id=utils.get_channel_id(peer_id), access_hash=0 - )] - ) - ) - else: - self.send( - functions.messages.GetChats( - id=[-peer_id] - ) + ) + ] ) + ) try: return self.storage.get_peer_by_id(peer_id) diff --git a/pyrogram/client/ext/utils.py b/pyrogram/client/ext/utils.py index e0a797e2..d89f83bb 100644 --- a/pyrogram/client/ext/utils.py +++ b/pyrogram/client/ext/utils.py @@ -18,10 +18,11 @@ import base64 import struct -from typing import Union, List +from typing import List +from typing import Union import pyrogram - +from pyrogram.api.types import PeerUser, PeerChat, PeerChannel from . import BaseClient from ...api import types @@ -62,23 +63,6 @@ def encode(s: bytes) -> str: return base64.urlsafe_b64encode(r).decode().rstrip("=") -def get_peer_id(input_peer) -> int: - return ( - input_peer.user_id if isinstance(input_peer, types.InputPeerUser) - else -input_peer.chat_id if isinstance(input_peer, types.InputPeerChat) - else int("-100" + str(input_peer.channel_id)) - ) - - -def get_input_peer(peer_id: int, access_hash: int): - return ( - types.InputPeerUser(user_id=peer_id, access_hash=access_hash) if peer_id > 0 - else types.InputPeerChannel(channel_id=int(str(peer_id)[4:]), access_hash=access_hash) - if (str(peer_id).startswith("-100") and access_hash) - else types.InputPeerChat(chat_id=-peer_id) - ) - - def get_offset_date(dialogs): for m in reversed(dialogs.messages): if isinstance(m, types.MessageEmpty): @@ -183,7 +167,7 @@ def parse_deleted_messages(client, update) -> List["pyrogram.Message"]: pyrogram.Message( message_id=message, chat=pyrogram.Chat( - id=int("-100" + str(channel_id)), + id=get_channel_id(channel_id), type="channel", client=client ) if channel_id is not None else None, @@ -203,3 +187,39 @@ def unpack_inline_message_id(inline_message_id: str) -> types.InputBotInlineMess id=r[1], access_hash=r[2] ) + + +MIN_CHANNEL_ID = -1002147483647 +MAX_CHANNEL_ID = -1000000000000 +MIN_CHAT_ID = -2147483647 +MAX_USER_ID = 2147483647 + + +def get_peer_id(peer: Union[PeerUser, PeerChat, PeerChannel]) -> int: + if isinstance(peer, PeerUser): + return peer.user_id + + if isinstance(peer, PeerChat): + return -peer.chat_id + + if isinstance(peer, PeerChannel): + return MAX_CHANNEL_ID - peer.channel_id + + raise ValueError("Peer type invalid: {}".format(peer)) + + +def get_type(peer_id: int) -> str: + if peer_id < 0: + if MIN_CHAT_ID <= peer_id: + return "chat" + + if MIN_CHANNEL_ID <= peer_id < MAX_CHANNEL_ID: + return "channel" + elif 0 < peer_id <= MAX_USER_ID: + return "user" + + raise ValueError("Peer id invalid: {}".format(peer_id)) + + +def get_channel_id(peer_id: int) -> int: + return MAX_CHANNEL_ID - peer_id diff --git a/pyrogram/client/methods/chats/get_chat.py b/pyrogram/client/methods/chats/get_chat.py index 48c5cc22..0773ce6c 100644 --- a/pyrogram/client/methods/chats/get_chat.py +++ b/pyrogram/client/methods/chats/get_chat.py @@ -20,7 +20,7 @@ from typing import Union import pyrogram from pyrogram.api import functions, types -from ...ext import BaseClient +from ...ext import BaseClient, utils class GetChat(BaseClient): @@ -70,7 +70,7 @@ class GetChat(BaseClient): chat_id = -r.chat.id if isinstance(r.chat, types.Channel): - chat_id = int("-100" + str(r.chat.id)) + chat_id = utils.get_channel_id(r.chat.id) peer = self.resolve_peer(chat_id) diff --git a/pyrogram/client/methods/chats/get_dialogs.py b/pyrogram/client/methods/chats/get_dialogs.py index f77ad30a..30078d57 100644 --- a/pyrogram/client/methods/chats/get_dialogs.py +++ b/pyrogram/client/methods/chats/get_dialogs.py @@ -23,7 +23,7 @@ from typing import List import pyrogram from pyrogram.api import functions, types from pyrogram.errors import FloodWait -from ...ext import BaseClient +from ...ext import BaseClient, utils log = logging.getLogger(__name__) @@ -100,10 +100,8 @@ class GetDialogs(BaseClient): chat_id = to_id.user_id else: chat_id = message.from_id - elif isinstance(to_id, types.PeerChat): - chat_id = -to_id.chat_id else: - chat_id = int("-100" + str(to_id.channel_id)) + chat_id = utils.get_peer_id(to_id) messages[chat_id] = pyrogram.Message._parse(self, message, users, chats) diff --git a/pyrogram/client/types/bots_and_keyboards/callback_query.py b/pyrogram/client/types/bots_and_keyboards/callback_query.py index d58865b2..9ba1804b 100644 --- a/pyrogram/client/types/bots_and_keyboards/callback_query.py +++ b/pyrogram/client/types/bots_and_keyboards/callback_query.py @@ -25,6 +25,7 @@ from pyrogram.api import types from ..object import Object from ..update import Update from ..user_and_chats import User +from ...ext import utils class CallbackQuery(Object, Update): @@ -90,16 +91,7 @@ class CallbackQuery(Object, Update): inline_message_id = None if isinstance(callback_query, types.UpdateBotCallbackQuery): - peer = callback_query.peer - - if isinstance(peer, types.PeerUser): - peer_id = peer.user_id - elif isinstance(peer, types.PeerChat): - peer_id = -peer.chat_id - else: - peer_id = int("-100" + str(peer.channel_id)) - - message = client.get_messages(peer_id, callback_query.msg_id) + message = client.get_messages(utils.get_peer_id(callback_query.peer), callback_query.msg_id) elif isinstance(callback_query, types.UpdateInlineBotCallbackQuery): inline_message_id = b64encode( pack( diff --git a/pyrogram/client/types/messages_and_media/message.py b/pyrogram/client/types/messages_and_media/message.py index a5105888..ed0088cd 100644 --- a/pyrogram/client/types/messages_and_media/message.py +++ b/pyrogram/client/types/messages_and_media/message.py @@ -31,7 +31,8 @@ from ..object import Object from ..update import Update from ..user_and_chats.chat import Chat from ..user_and_chats.user import User -from ...parser import utils, Parser +from ...ext import utils +from ...parser import utils as parser_utils, Parser class Str(str): @@ -54,7 +55,7 @@ class Str(str): return Parser.unparse(self, self.entities, True) def __getitem__(self, item): - return utils.remove_surrogates(utils.add_surrogates(self)[item]) + return parser_utils.remove_surrogates(parser_utils.add_surrogates(self)[item]) class Message(Object, Update): @@ -446,7 +447,7 @@ class Message(Object, Update): new_chat_title=new_chat_title, new_chat_photo=new_chat_photo, delete_chat_photo=delete_chat_photo, - migrate_to_chat_id=int("-100" + str(migrate_to_chat_id)) if migrate_to_chat_id else None, + migrate_to_chat_id=utils.get_channel_id(migrate_to_chat_id) if migrate_to_chat_id else None, migrate_from_chat_id=-migrate_from_chat_id if migrate_from_chat_id else None, group_chat_created=group_chat_created, channel_chat_created=channel_chat_created, diff --git a/pyrogram/client/types/user_and_chats/chat.py b/pyrogram/client/types/user_and_chats/chat.py index cac5d0c7..396831cc 100644 --- a/pyrogram/client/types/user_and_chats/chat.py +++ b/pyrogram/client/types/user_and_chats/chat.py @@ -23,6 +23,7 @@ from pyrogram.api import types from .chat_permissions import ChatPermissions from .chat_photo import ChatPhoto from ..object import Object +from ...ext import utils class Chat(Object): @@ -180,7 +181,7 @@ class Chat(Object): @staticmethod def _parse_channel_chat(client, channel: types.Channel) -> "Chat": - peer_id = int("-100" + str(channel.id)) + peer_id = utils.get_channel_id(channel.id) return Chat( id=peer_id, @@ -672,7 +673,7 @@ class Chat(Object): can_pin_messages=can_pin_messages, can_promote_members=can_promote_members ) - + def join(self): """Bound method *join* of :obj:`Chat`. diff --git a/pyrogram/client/types/user_and_chats/dialog.py b/pyrogram/client/types/user_and_chats/dialog.py index 4ea82184..a78e501b 100644 --- a/pyrogram/client/types/user_and_chats/dialog.py +++ b/pyrogram/client/types/user_and_chats/dialog.py @@ -21,6 +21,7 @@ import pyrogram from pyrogram.api import types from ..object import Object from ..user_and_chats import Chat +from ...ext import utils class Dialog(Object): @@ -70,18 +71,9 @@ class Dialog(Object): @staticmethod def _parse(client, dialog: types.Dialog, messages, users, chats) -> "Dialog": - chat_id = dialog.peer - - if isinstance(chat_id, types.PeerUser): - chat_id = chat_id.user_id - elif isinstance(chat_id, types.PeerChat): - chat_id = -chat_id.chat_id - else: - chat_id = int("-100" + str(chat_id.channel_id)) - return Dialog( chat=Chat._parse_dialog(client, dialog.peer, users, chats), - top_message=messages.get(chat_id), + top_message=messages.get(utils.get_peer_id(dialog.peer)), unread_messages_count=dialog.unread_count, unread_mentions_count=dialog.unread_mentions_count, unread_mark=dialog.unread_mark,