deluge/deluge/transfer.py

# -*- coding: utf-8 -*-
#
# Copyright (C) 2012 Bro <bro.development@gmail.com>
#
# This file is part of Deluge and is licensed under GNU General Public License 3.0, or later, with
# the additional special exception to link portions of this program with the OpenSSL library.
# See LICENSE for more details.
#

from __future__ import unicode_literals

import logging
import struct
import zlib

from twisted.internet.protocol import Protocol

try:
    import rencode  # pylint: disable=useless-suppression,relative-import
except ImportError:
    import deluge.rencode as rencode

log = logging.getLogger(__name__)

MESSAGE_HEADER_SIZE = 5


class DelugeTransferProtocol(Protocol, object):
    """
    Data messages are transfered using very a simple protocol.
    Data messages are transfered with a header containing
    the length of the data to be transfered (payload).

    """

    def __init__(self):
        self._buffer = b''  # TODO: Look into using bytearray instead of byte string.
        self._message_length = 0
        self._bytes_received = 0
        self._bytes_sent = 0

    def transfer_message(self, data):
        """
        Transfer the data.

        The data will be serialized and compressed before being sent.
        First a header is sent - containing the length of the compressed payload
        to come as a signed integer. After the header, the payload is transfered.

        :param data: data to be transfered in a data structure serializable by rencode.

        """
        compressed = zlib.compress(rencode.dumps(data))
        size_data = len(compressed)
        # Store length as a signed integer (using 4 bytes). "!" denotes network byte order.
        payload_len = struct.pack('!i', size_data)
        header = b'D' + payload_len
        self._bytes_sent += len(header) + len(compressed)
        self.transport.write(header)
        self.transport.write(compressed)

    def dataReceived(self, data):  # NOQA: N802
        """
        This method is called whenever data is received.

        :param data: a message as transfered by transfer_message, or a part of such
                     a messsage.

        Global variables:
            _buffer         - contains the data received
            _message_length - the length of the payload of the current message.

        """
        self._buffer += data
        self._bytes_received += len(data)

        while len(self._buffer) >= MESSAGE_HEADER_SIZE:
            if self._message_length == 0:
                self._handle_new_message()
            # We have a complete packet
            if len(self._buffer) >= self._message_length:
                self._handle_complete_message(self._buffer[: self._message_length])
                # Remove message data from buffer
                self._buffer = self._buffer[self._message_length :]
                self._message_length = 0
            else:
                break

    def _handle_new_message(self):
        """
        Handle the start of a new message. This method is called only when the
        beginning of the buffer contains data from a new message (i.e. the header).

        """
        try:
            # Read the first bytes of the message (MESSAGE_HEADER_SIZE bytes)
            header = self._buffer[:MESSAGE_HEADER_SIZE]
            payload_len = header[1:MESSAGE_HEADER_SIZE]
            if header[0:1] != b'D':
                raise Exception(
                    'Invalid header format. First byte is %d' % ord(header[0:1])
                )
            # Extract the length stored as a signed integer (using 4 bytes)
            self._message_length = struct.unpack('!i', payload_len)[0]
            if self._message_length < 0:
                raise Exception('Message length is negative: %d' % self._message_length)
            # Remove the header from the buffer
            self._buffer = self._buffer[MESSAGE_HEADER_SIZE:]
        except Exception as ex:
            log.warning('Error occurred when parsing message header: %s.', ex)
            log.warning(
                'This version of Deluge cannot communicate with the sender of this data.'
            )
            self._message_length = 0
            self._buffer = b''

    def _handle_complete_message(self, data):
        """
        Handles a complete message as it is transfered on the network.

        :param data: a zlib compressed string encoded with rencode.

        """
        try:
            self.message_received(
                rencode.loads(zlib.decompress(data), decode_utf8=True)
            )
        except Exception as ex:
            log.warning(
                'Failed to decompress (%d bytes) and load serialized data with rencode: %s',
                len(data),
                ex,
            )

    def get_bytes_recv(self):
        """
        Returns the number of bytes received.

        :returns: the number of bytes received
        :rtype: int

        """
        return self._bytes_received

    def get_bytes_sent(self):
        """
        Returns the number of bytes sent.

        :returns: the number of bytes sent
        :rtype: int

        """
        return self._bytes_sent

    def message_received(self, message):
        """Override this method to receive the complete message"""
        pass