peerwiretranslator.py

"""
The PeerWireTranslator is responsible for translating a stream of bytes into
higher level messages and vice versa related to the peer wire protocol between
BitTorrent peers.  The bytes are exchanged with a readerwriter and the higher
level messages are exchanged with a receiver.  The creator of the
PeerWireTranslator must provide it with a readerwriter upon creation or
subsequently using the set_readerwriter() method.  When it no longer needs the
PeerWireTranslator, it should call the unset_readerwriter() method so that all
references to the PeerWireTranslator can be released and it can be properly
garbage collected.

On the receiver side, when told to transmit one of the peer wire protocol
messages, the PeerWireTranslator generates a set of bytes to send to the
readerwriter.  When the readerwriter presents the PeerWireTranslator with a set
of bytes that represent one of the peer wire protocol messages, the translator
notifies the receiver.  It also notifies the receiver when a connection is
lost.

A receiver must implement the following methods: rx_keep_alive(), rx_choke(),
rx_unchoke(), rx_interested(), rx_not_interested, rx_bitfield(), rx_have(),
rx_request(), rx_piece() and rx_cancel() and connection_lost().

On the readerwriter side, when incoming bytes are available, the readerwriter
asks the PeerWireTranslator for a buffer to put them into and after it has
done that, it notifies the translator that bytes have been received.  This
scheme allows the PeerWireTranslator to get enough bytes to translate the
next part of the message expected.  It also allows an entire block of incoming
data to be placed in one buffer even if it is received over multiple socket
reads.

A readerwriter must implement set_receiver(), unset_receiver() and tx_bytes()
"""

import logging
import struct
from bitstring import BitArray

logger = logging.getLogger('bt.bttranslator')

_LENGTH_LEN = 4

_MSG_CHOKE = 0
_MSG_UNCHOKE = 1
_MSG_INTERESTED = 2
_MSG_NOT_INTERESTED = 3
_MSG_HAVE = 4
_MSG_BITFIELD = 5
_MSG_REQUEST = 6
_MSG_PIECE = 7
_MSG_CANCEL = 8


class PeerWireTranslator(object):
    class _States(object):
        Length, Message = range(2)

    def __init__(self, receiver=None, readerwriter=None):
        self._length_buf = bytearray(_LENGTH_LEN)
        self._length_view = memoryview(self._length_buf)
        self._length_state_setup()

        if receiver is not None:
            self.set_receiver(receiver)
        else:
            self._receiver = None

        if readerwriter is not None:
            self.set_readerwriter(readerwriter)
        else:
            self._readerwriter = None

        self._rx_functions = {_MSG_CHOKE: self.rx_choke,
                              _MSG_UNCHOKE: self.rx_unchoke,
                              _MSG_INTERESTED: self.rx_interested,
                              _MSG_NOT_INTERESTED: self.rx_not_interested,
                              _MSG_HAVE: self.rx_have,
                              _MSG_BITFIELD: self.rx_bitfield,
                              _MSG_REQUEST: self.rx_request,
                              _MSG_PIECE: self.rx_piece,
                              _MSG_CANCEL: self.rx_cancel}

    def _length_state_setup(self):
        self._rx_state = self._States.Length
        self._bytes_needed = _LENGTH_LEN
        self._bytes_received = 0
        self._current_buf = self._length_buf
        self._current_view = self._length_view

    def set_receiver(self, receiver):
        self._receiver = receiver

    def unset_receiver(self):
        self._receiver = None

    def set_readerwriter(self, readerwriter):
        self._readerwriter = readerwriter
        self._readerwriter.set_receiver(self)

    def unset_readerwriter(self):
        self._readerwriter.unset_receiver()
        self._readerwriter = None

    def get_rx_buffer(self):
        return self._current_view[self._bytes_received:], self._bytes_needed

    def rx_bytes(self, n):
        self._bytes_received += n
        self._bytes_needed -= n

        if self._bytes_needed == 0:
            if self._rx_state == self._States.Length:
                (length,) = struct.unpack('>i', buffer(self._length_buf))
                if length == 0:
                    self.rx_keep_alive()
                    self._length_state_setup()
                else:
                    self._rx_state = self._States.Message
                    self._bytes_needed = length
                    self._bytes_received = 0

                    self._current_buf = bytearray(length)
                    self._current_view = memoryview(self._current_buf)
            else:
                (message_id,) = struct.unpack("B",
                                              buffer(self._current_buf[0:1]))

                try:
                    self._rx_functions[message_id]()
                except KeyError:
                    logger.debug("Received message with invalid msg id: {}"
                                 .format(message_id))

                self._length_state_setup()

    def rx_keep_alive(self):
        if self._receiver is not None:
            self._receiver.rx_keep_alive()

    def rx_choke(self):
        if self._receiver is not None:
            self._receiver.rx_choke()

    def rx_unchoke(self):
        if self._receiver is not None:
            self._receiver.rx_unchoke()

    def rx_interested(self):
        if self._receiver is not None:
            self._receiver.rx_interested()

    def rx_not_interested(self):
        if self._receiver is not None:
            self._receiver.rx_not_interested()

    def rx_have(self):
        if self._receiver is not None:
            (index,) = struct.unpack(">I", buffer(self._current_buf[1:5]))
            self._receiver.rx_have(index)

    def rx_bitfield(self):
        if self._receiver is not None:
            bits = BitArray(bytes=self._current_buf[1:self._bytes_received])
            self._receiver.rx_bitfield(bits)

    def rx_request(self):
        if self._receiver is not None:
            buf = buffer(self._current_buf[1:])
            index, begin, length, = struct.unpack(">3I", buf)
            self._receiver.rx_request(index, begin, length)

    def rx_piece(self):
        if self._receiver is not None:
            index, begin, = struct.unpack(">2I",
                                          buffer(self._current_buf[1:9]))
            self._receiver.rx_piece(index, begin,
                                    buffer(self._current_buf[9:]))

    def rx_cancel(self):
        if self._receiver is not None:
            buf = buffer(self._current_buf[1:])
            index, begin, length, = struct.unpack(">3I", buf)
            self._receiver.rx_cancel(index, begin, length)

    def tx_keep_alive(self):
        if self._readerwriter is not None:
            self._readerwriter.tx_bytes(struct.pack('B', 0))

    def tx_choke(self):
        if self._readerwriter is not None:
            self._readerwriter.tx_bytes(struct.pack('>IB', 1, _MSG_CHOKE))

    def tx_unchoke(self):
        if self._readerwriter is not None:
            self._readerwriter.tx_bytes(struct.pack('>IB', 1, _MSG_UNCHOKE))

    def tx_interested(self):
        if self._readerwriter is not None:
            self._readerwriter.tx_bytes(struct.pack('>IB', 1, _MSG_INTERESTED))

    def tx_not_interested(self):
        if self._readerwriter is not None:
            self._readerwriter.tx_bytes(struct.pack('>IB', 1,
                                                    _MSG_NOT_INTERESTED))

    def tx_have(self, index):
        if self._readerwriter is not None:
            self._readerwriter.tx_bytes(struct.pack('>IBI', 5,
                                                    _MSG_HAVE, index))

    def tx_bitfield(self, bits):
        if self._readerwriter is not None:
            bitfield = bits.tobytes()
            length = len(bitfield)
            self._readerwriter.tx_bytes(struct.pack('>IB{}s'.format(length),
                                                    1+length, _MSG_BITFIELD,
                                                    bitfield))

    def tx_request(self, index, begin, length):
        if self._readerwriter is not None:
            self._readerwriter.tx_bytes(struct.pack('>IB3I', 13, _MSG_REQUEST,
                                                    index, begin, length))

    def tx_piece(self, index, begin, block):
        if self._readerwriter is not None:
            length = len(block)
            self._readerwriter.tx_bytes(struct.pack('>IB2I{}s'.format(length),
                                                    9+length, _MSG_PIECE,
                                                    index, begin, block))

    def tx_cancel(self, index, begin, length):
        if self._readerwriter is not None:
            self._readerwriter.tx_bytes(struct.pack('>IB3I', 13, _MSG_CANCEL,
                                                    index, begin, length))

    def connection_lost(self):
        if self._receiver is not None:
            self._receiver.connection_lost()