<commit_before>"""
This module contains the class used to communicate with Telegram's servers
encrypting every packet, and relies on a valid AuthKey in the used Session.
"""
import gzip
import logging
import struct
from threading import Lock

from .. import helpers as utils
from ..crypto import AES
from ..errors import (
    BadMessageError, InvalidChecksumError, BrokenAuthKeyError,
    rpc_message_to_error
)
from ..extensions import BinaryReader
from ..tl import TLMessage, MessageContainer, GzipPacked
from ..tl.all_tlobjects import tlobjects
from ..tl.types import (
    MsgsAck, Pong, BadServerSalt, BadMsgNotification,
    MsgNewDetailedInfo, NewSessionCreated, MsgDetailedInfo
)
from ..tl.functions.auth import LogOutRequest

__log__ = logging.getLogger(__name__)


class MtProtoSender:
    """MTProto Mobile Protocol sender
       (https://core.telegram.org/mtproto/description).

       Note that this class is not thread-safe, and calling send/receive
       from two or more threads at the same time is undefined behaviour.
       Rationale: a new connection should be spawned to send/receive requests
                  in parallel, so thread-safety (hence locking) isn't needed.
    """

    def __init__(self, session, connection):
        """
        Initializes a new MTProto sender.

        :param session:
            the Session to be used with this sender. Must contain the IP and
            port of the server, salt, ID, and AuthKey,
        :param connection:
            the Connection to be used.
        """
        self.session = session
        self.connection = connection

        # Message IDs that need confirmation
        self._need_confirmation = set()

        # Requests (as msg_id: Message) sent waiting to be received
        self._pending_receive = {}

        # Multithreading
        self._send_lock = Lock()

    def connect(self):
        """Connects to the server."""
        self.connection.connect(self.session.server_address, self.session.port)

    def is_connected(self):
        """
        Determines whether the sender is connected or not.

        :return: true if the sender is connected.
        """
        return self.connection.is_connected()

    def disconnect(self):
        """Disconnects from the server."""
        self.connection.close()
        self._need_confirmation.clear()
        self._clear_all_pending()

    # region Send and receive

    def send(self, *requests):
        """
        Sends the specified TLObject(s) (which must be requests),
        and acknowledging any message which needed confirmation.

        :param requests: the requests to be sent.
        """
        # Finally send our packed request(s)
        messages = [TLMessage(self.session, r) for r in requests]
        self._pending_receive.update({m.msg_id: m for m in messages})

        # Pack everything in the same container if we need to send AckRequests
        if self._need_confirmation:
            messages.append(
                TLMessage(self.session, MsgsAck(list(self._need_confirmation)))
            )
            self._need_confirmation.clear()

        if len(messages) == 1:
            message = messages[0]
        else:
            message = TLMessage(self.session, MessageContainer(messages))
            # On bad_msg_salt errors, Telegram will reply with the ID of
            # the container and not the requests it contains, so in case
            # this happens we need to know to which container they belong.
            for m in messages:
                m.container_msg_id = message.msg_id

        self._send_message(message)

    def _send_acknowledge(self, msg_id):
        """Sends a message acknowledge for the given msg_id."""
        self._send_message(TLMessage(self.session, MsgsAck([msg_id])))

    def receive(self, update_state):
        """
        Receives a single message from the connected endpoint.

        This method returns nothing, and will only affect other parts
        of the MtProtoSender such as the updates callback being fired
        or a pending request being confirmed.

        Any unhandled object (likely updates) will be passed to
        update_state.process(TLObject).

        :param update_state:
            the UpdateState that will process all the received
            Update and Updates objects.
        """
        try:
            body = self.connection.recv()
        except (BufferError, InvalidChecksumError):
            # TODO BufferError, we should spot the cause...
            # "No more bytes left"; something wrong happened, clear
            # everything to be on the safe side, or:
            #
            # "This packet should be skipped"; since this may have
            # been a result for a request, invalidate every request
            # and just re-invoke them to avoid problems
            __log__.exception('Error while receiving server response. '
                              '%d pending request(s) will be ignored',
                              len(self._pending_receive))
            self._clear_all_pending()
            return

        message, remote_msg_id, remote_seq = self._decode_msg(body)
        with BinaryReader(message) as reader:
            self._process_msg(remote_msg_id, remote_seq, reader, update_state)

    # endregion

    # region Low level processing

    def _send_message(self, message):
        """
        Sends the given encrypted through the network.

        :param message: the TLMessage to be sent.
        """
        with self._send_lock:
            self.connection.send(utils.pack_message(self.session, message))

    def _decode_msg(self, body):
        """
        Decodes the body of the payload received from the network.

        :param body: the body to be decoded.
        :return: a tuple of (decoded message, remote message id, remote seq).
        """
        if len(body) < 8:
            if body == b'l\xfe\xff\xff':
                raise BrokenAuthKeyError()
            else:
                raise BufferError("Can't decode packet ({})".format(body))

        with BinaryReader(body) as reader:
            return utils.unpack_message(self.session, reader)

    def _process_msg(self, msg_id, sequence, reader, state):
        """
        Processes the message read from the network inside reader.

        :param msg_id: the ID of the message.
        :param sequence: the sequence of the message.
        :param reader: the BinaryReader that contains the message.
        :param state: the current UpdateState.
        :return: true if the message was handled correctly, false otherwise.
        """
        # TODO Check salt, session_id and sequence_number
        self._need_confirmation.add(msg_id)

        code = reader.read_int(signed=False)
        reader.seek(-4)

        # These are a bit of special case, not yet generated by the code gen
        if code == 0xf35c6d01:  # rpc_result, (response of an RPC call)
            __log__.debug('Processing Remote Procedure Call result')
            return self._handle_rpc_result(msg_id, sequence, reader)

        if code == MessageContainer.CONSTRUCTOR_ID:
            __log__.debug('Processing container result')
            return self._handle_container(msg_id, sequence, reader, state)

        if code == GzipPacked.CONSTRUCTOR_ID:
            __log__.debug('Processing gzipped result')
            return self._handle_gzip_packed(msg_id, sequence, reader, state)

        if code not in tlobjects:
            __log__.warning(
                'Unknown message with ID %d, data left in the buffer %s',
                hex(code), repr(reader.get_bytes()[reader.tell_position():])
            )
            return False

        obj = reader.tgread_object()
        __log__.debug('Processing %s result', type(obj).__name__)

        if isinstance(obj, Pong):
            return self._handle_pong(msg_id, sequence, obj)

        if isinstance(obj, BadServerSalt):
            return self._handle_bad_server_salt(msg_id, sequence, obj)

        if isinstance(obj, BadMsgNotification):
            return self._handle_bad_msg_notification(msg_id, sequence, obj)

        if isinstance(obj, MsgDetailedInfo):
            return self._handle_msg_detailed_info(msg_id, sequence, obj)

        if isinstance(obj, MsgNewDetailedInfo):
            return self._handle_msg_new_detailed_info(msg_id, sequence, obj)

        if isinstance(obj, NewSessionCreated):
            return self._handle_new_session_created(msg_id, sequence, obj)

        if isinstance(obj, MsgsAck):  # may handle the request we wanted
            # Ignore every ack request *unless* when logging out, when it's
            # when it seems to only make sense. We also need to set a non-None
            # result since Telegram doesn't send the response for these.
            for msg_id in obj.msg_ids:
                r = self._pop_request_of_type(msg_id, LogOutRequest)
                if r:
                    r.result = True  # Telegram won't send this value
                    r.confirm_received.set()

            return True

        # If the object isn't any of the above, then it should be an Update.
        self.session.process_entities(obj)
        if state:
            state.process(obj)

        return True

    # endregion

    # region Message handling

    def _pop_request(self, msg_id):
        """
        Pops a pending **request** from self._pending_receive.

        :param msg_id: the ID of the message that belongs to the request.
        :return: the request, or None if it wasn't found.
        """
        message = self._pending_receive.pop(msg_id, None)
        if message:<commit_msg>Add missing FutureSalts response special case (#81)<commit_after>"""
This module contains the class used to communicate with Telegram's servers
encrypting every packet, and relies on a valid AuthKey in the used Session.
"""
import gzip
import logging
from threading import Lock

from .. import helpers as utils
from ..errors import (
    BadMessageError, InvalidChecksumError, BrokenAuthKeyError,
    rpc_message_to_error
)
from ..extensions import BinaryReader
from ..tl import TLMessage, MessageContainer, GzipPacked
from ..tl.all_tlobjects import tlobjects
from ..tl.functions.auth import LogOutRequest
from ..tl.types import (
    MsgsAck, Pong, BadServerSalt, BadMsgNotification, FutureSalts,
    MsgNewDetailedInfo, NewSessionCreated, MsgDetailedInfo
)

__log__ = logging.getLogger(__name__)


class MtProtoSender:
    """MTProto Mobile Protocol sender
       (https://core.telegram.org/mtproto/description).

       Note that this class is not thread-safe, and calling send/receive
       from two or more threads at the same time is undefined behaviour.
       Rationale: a new connection should be spawned to send/receive requests
                  in parallel, so thread-safety (hence locking) isn't needed.
    """

    def __init__(self, session, connection):
        """
        Initializes a new MTProto sender.

        :param session:
            the Session to be used with this sender. Must contain the IP and
            port of the server, salt, ID, and AuthKey,
        :param connection:
            the Connection to be used.
        """
        self.session = session
        self.connection = connection

        # Message IDs that need confirmation
        self._need_confirmation = set()

        # Requests (as msg_id: Message) sent waiting to be received
        self._pending_receive = {}

        # Multithreading
        self._send_lock = Lock()

    def connect(self):
        """Connects to the server."""
        self.connection.connect(self.session.server_address, self.session.port)

    def is_connected(self):
        """
        Determines whether the sender is connected or not.

        :return: true if the sender is connected.
        """
        return self.connection.is_connected()

    def disconnect(self):
        """Disconnects from the server."""
        self.connection.close()
        self._need_confirmation.clear()
        self._clear_all_pending()

    # region Send and receive

    def send(self, *requests):
        """
        Sends the specified TLObject(s) (which must be requests),
        and acknowledging any message which needed confirmation.

        :param requests: the requests to be sent.
        """
        # Finally send our packed request(s)
        messages = [TLMessage(self.session, r) for r in requests]
        self._pending_receive.update({m.msg_id: m for m in messages})

        # Pack everything in the same container if we need to send AckRequests
        if self._need_confirmation:
            messages.append(
                TLMessage(self.session, MsgsAck(list(self._need_confirmation)))
            )
            self._need_confirmation.clear()

        if len(messages) == 1:
            message = messages[0]
        else:
            message = TLMessage(self.session, MessageContainer(messages))
            # On bad_msg_salt errors, Telegram will reply with the ID of
            # the container and not the requests it contains, so in case
            # this happens we need to know to which container they belong.
            for m in messages:
                m.container_msg_id = message.msg_id

        self._send_message(message)

    def _send_acknowledge(self, msg_id):
        """Sends a message acknowledge for the given msg_id."""
        self._send_message(TLMessage(self.session, MsgsAck([msg_id])))

    def receive(self, update_state):
        """
        Receives a single message from the connected endpoint.

        This method returns nothing, and will only affect other parts
        of the MtProtoSender such as the updates callback being fired
        or a pending request being confirmed.

        Any unhandled object (likely updates) will be passed to
        update_state.process(TLObject).

        :param update_state:
            the UpdateState that will process all the received
            Update and Updates objects.
        """
        try:
            body = self.connection.recv()
        except (BufferError, InvalidChecksumError):
            # TODO BufferError, we should spot the cause...
            # "No more bytes left"; something wrong happened, clear
            # everything to be on the safe side, or:
            #
            # "This packet should be skipped"; since this may have
            # been a result for a request, invalidate every request
            # and just re-invoke them to avoid problems
            __log__.exception('Error while receiving server response. '
                              '%d pending request(s) will be ignored',
                              len(self._pending_receive))
            self._clear_all_pending()
            return

        message, remote_msg_id, remote_seq = self._decode_msg(body)
        with BinaryReader(message) as reader:
            self._process_msg(remote_msg_id, remote_seq, reader, update_state)

    # endregion

    # region Low level processing

    def _send_message(self, message):
        """
        Sends the given encrypted through the network.

        :param message: the TLMessage to be sent.
        """
        with self._send_lock:
            self.connection.send(utils.pack_message(self.session, message))

    def _decode_msg(self, body):
        """
        Decodes the body of the payload received from the network.

        :param body: the body to be decoded.
        :return: a tuple of (decoded message, remote message id, remote seq).
        """
        if len(body) < 8:
            if body == b'l\xfe\xff\xff':
                raise BrokenAuthKeyError()
            else:
                raise BufferError("Can't decode packet ({})".format(body))

        with BinaryReader(body) as reader:
            return utils.unpack_message(self.session, reader)

    def _process_msg(self, msg_id, sequence, reader, state):
        """
        Processes the message read from the network inside reader.

        :param msg_id: the ID of the message.
        :param sequence: the sequence of the message.
        :param reader: the BinaryReader that contains the message.
        :param state: the current UpdateState.
        :return: true if the message was handled correctly, false otherwise.
        """
        # TODO Check salt, session_id and sequence_number
        self._need_confirmation.add(msg_id)

        code = reader.read_int(signed=False)
        reader.seek(-4)

        # These are a bit of special case, not yet generated by the code gen
        if code == 0xf35c6d01:  # rpc_result, (response of an RPC call)
            __log__.debug('Processing Remote Procedure Call result')
            return self._handle_rpc_result(msg_id, sequence, reader)

        if code == MessageContainer.CONSTRUCTOR_ID:
            __log__.debug('Processing container result')
            return self._handle_container(msg_id, sequence, reader, state)

        if code == GzipPacked.CONSTRUCTOR_ID:
            __log__.debug('Processing gzipped result')
            return self._handle_gzip_packed(msg_id, sequence, reader, state)

        if code not in tlobjects:
            __log__.warning(
                'Unknown message with ID %d, data left in the buffer %s',
                hex(code), repr(reader.get_bytes()[reader.tell_position():])
            )
            return False

        obj = reader.tgread_object()
        __log__.debug('Processing %s result', type(obj).__name__)

        if isinstance(obj, Pong):
            return self._handle_pong(msg_id, sequence, obj)

        if isinstance(obj, BadServerSalt):
            return self._handle_bad_server_salt(msg_id, sequence, obj)

        if isinstance(obj, BadMsgNotification):
            return self._handle_bad_msg_notification(msg_id, sequence, obj)

        if isinstance(obj, MsgDetailedInfo):
            return self._handle_msg_detailed_info(msg_id, sequence, obj)

        if isinstance(obj, MsgNewDetailedInfo):
            return self._handle_msg_new_detailed_info(msg_id, sequence, obj)

        if isinstance(obj, NewSessionCreated):
            return self._handle_new_session_created(msg_id, sequence, obj)

        if isinstance(obj, MsgsAck):  # may handle the request we wanted
            # Ignore every ack request *unless* when logging out, when it's
            # when it seems to only make sense. We also need to set a non-None
            # result since Telegram doesn't send the response for these.
            for msg_id in obj.msg_ids:
                r = self._pop_request_of_type(msg_id, LogOutRequest)
                if r:
                    r.result = True  # Telegram won't send this value
                    r.confirm_received.set()

            return True

        if isinstance(obj, FutureSalts):
            r = self._pop_request(obj.req_msg_id)
            if r:
                r.result = obj
                r.confirm_received.set()

        # If the object isn't any of the above, then it should be an Update.
        self.session.process_entities(obj)
        if state:
            state.process(obj)

        return True

    # endregion

    # region Message handling

    def _pop_request(self, msg_id):
        """
        Pops a pending **request** from self._pending_receive.

        :param msg_id: the ID of the message that belongs to the request.
        :return: the request, or None if it wasn't found.
        """
        message = self._pending_receive.pop(msg_id, None)
        if message:
{'fraction_non_alphanumeric': 0.07507692307692308, 'fraction_numerical': 0.0011282051282051281, 'mean_word_length': 2.2523348899266176, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2472655', 'n_tokens_mistral': 5184, 'n_tokens_neox': 4967, 'n_words': 1764}