Exemplo n.º 1
0
    def test_write(self):
        string1 = "test string"
        string2 = "second test string"

        w = SocketReaderWriter(wsock)
        w.tx_bytes(string1)
        
        select.select([],[wsock],[])
        w.write_event()
    
        # make sure unregister_for_write_events called
        select.select([rsock],[],[])
        data = rsock.recv(65535) 
        assert data == string1 

        w.tx_bytes(string1)
        w.tx_bytes(string2)

        select.select([],[wsock],[])
        w.write_event()

        select.select([rsock],[],[])
        data = rsock.recv(len(string1)) 
        assert data == string1

        select.select([rsock],[],[])
        data = rsock.recv(len(string2)) 
        assert data == string2
Exemplo n.º 2
0
    def test_read(self):
        r = SocketReaderWriter(rsock)
        self.buffer = bytearray(100)
        self.view = memoryview(self.buffer)

        rcvr = Mock()
        rcvr.get_rx_buffer = Mock(return_value = (self.view, 100))

        string1 = "test string"
        string2 = "second test string"

        # bytes should be read only when the receiver is set
        wsock.send(string1)
        _, _, _ = select.select([rsock],[],[])
        r.read_event()
        
        r.set_receiver(rcvr)

        wsock.send(string2)

        select.select([rsock],[],[])
        r.read_event()

        total_length = len(string1)+len(string2)
        rcvr.rx_bytes.assert_called_with(total_length)
        assert self.view[0:total_length].tobytes() == string1+string2 

        # after the receiver is unset, make sure it isn't called for 
        # incoming bytes
        r.unset_receiver()
        wsock.send(string2)

        call_count = rcvr.rx_bytes.call_count 
        select.select([rsock],[],[])
        r.read_event()
        assert rcvr.rx_bytes.call_count == call_count
    
        # cleanup by reading the bytes from the socket
        r.set_receiver(rcvr)
        r.read_event()
Exemplo n.º 3
0
 def test_stream(self):
     r = SocketReaderWriter(rsock)
     assert r.stream() == rsock 
Exemplo n.º 4
0
    def test_partial_send(self):
        mock_sock = Mock()
        mock_sock.send = Mock(return_value = 3)

        string1 = "abcdefghi"
        w = SocketReaderWriter(mock_sock)
        w.tx_bytes(string1)

        w.write_event() 
        assert mock_sock.send.called_with(string1)

        w.write_event() 
        assert mock_sock.send.called_with(string1[3:])
    
        w.write_event() 
        assert mock_sock.send.called_with(string1[6:])

        call_count = mock_sock.send.call_count
        w.write_event()
        assert mock_sock.send.call_count == call_count
Exemplo n.º 5
0
    def _setup_handshake_translator(self):
        self._socketreaderwriter = SocketReaderWriter(self._socket)

        self._translator = HandshakeTranslator()
        self._translator.set_readerwriter(self._socketreaderwriter)
        self._translator.set_receiver(self)
Exemplo n.º 6
0
class PeerProxy(object):
    class _States(object):
        (Awaiting_Handshake, Awaiting_Connection, Handshake_Initiated,
         Bitfield_Allowed, Peer_to_Peer, Disconnected) = range(6)

    def __init__(self, client, peer_id, addr, socket=None, info_hash=None):
        self._client = client
        self._socket = socket
        self._info_hash = info_hash
        self._peer_id = peer_id
        self._addr = addr

        self._choked = True
        self._interested = False
        self._peer_choked = True
        self._peer_interested = False

        if len(peer_id) != 20:
            raise ValueError("Peer id must be 20 bytes long")

        if socket is None:
            if info_hash is None or len(info_hash) != 20:
                raise ValueError("Info hash must be a 20 byte value")

            self._translator = None
            self._socketreaderwriter = None
            Connector().connect(addr, self)

            self._state = self._States.Awaiting_Connection
        else:
            self._setup_handshake_translator()
            self._state = self._States.Awaiting_Handshake

    def _setup_handshake_translator(self):
        self._socketreaderwriter = SocketReaderWriter(self._socket)

        self._translator = HandshakeTranslator()
        self._translator.set_readerwriter(self._socketreaderwriter)
        self._translator.set_receiver(self)

    def _drop_connection(self, notify_client=True):
        if self._translator is not None:
            self._translator.unset_receiver()
            self._translator.unset_readerwriter()
            self._translator = None

        if self._socketreaderwriter:
            self._socketreaderwriter.stop()

        if self._socket is not None:
            self._socket.close()

        self._state = self._States.Disconnected

        if notify_client:
            self._client.peer_unconnected(self)

    def _valid_rx_state(self):
        if self._state != self._States.Peer_to_Peer:
            if self._state == self._States.Bitfield_Allowed:
                self._state = self._States.Peer_to_Peer
            else:
                if self._state != self._States.Disconnected:
                    self._drop_connection()
                return False
        return True

    def _valid_tx_state(self):
        if self._state != self._States.Peer_to_Peer:
            if self._state == self._States.Bitfield_Allowed:
                self._state = self._States.Peer_to_Peer
            else:
                return False
        return True

    def addr(self):
        return self._addr

    def is_interested(self):
        return self._interested

    def is_choked(self):
        return self._choked

    def is_peer_choked(self):
        return self._peer_choked

    def is_peer_interested(self):
        return self._peer_interested

    # Connector callbacks

    def connection_complete(self, addr, socket):
        self._socket = socket
        self._setup_handshake_translator()

        self._translator.tx_handshake(0, self._info_hash, self._peer_id)
        self._state = self._States.Handshake_Initiated

    def connection_failed(self, addr):
        self._client.peer_unconnected(self)

    # Translator callbacks

    def connection_lost(self):
        self._drop_connection()

    # HandshakeTranslator callbacks

    def rx_handshake(self, reserved, info_hash, peer_id):
        if self._state == self._States.Handshake_Initiated:
            if info_hash != self._info_hash:
                self._drop_connection()
            else:
                self._translator.unset_receiver()
                self._translator.unset_readerwriter()

                self._translator = PeerWireTranslator()
                self._translator.set_readerwriter(self._socketreaderwriter)
                self._translator.set_receiver(self)
                self._state = self._States.Bitfield_Allowed

                self._translator.tx_bitfield(self._client.get_bitfield())

    def rx_non_handshake(self):
        self._drop_connection()

    # PeerWireTranslator callbacks

    def rx_bitfield(self, bitfield):
        if self._state == self._States.Bitfield_Allowed:
            self._state = self._States.Peer_to_Peer
            self._client.peer_bitfield(self, bitfield)
        else:
            self._drop_connection()

    def rx_keep_alive(self):
        pass

    def rx_choke(self):
        if self._valid_rx_state():
            self._peer_choked = True
            self._client.peer_choked(self)

    def rx_unchoke(self):
        if self._valid_rx_state():
            self._peer_choked = False
            self._client.peer_unchoked(self)

    def rx_interested(self):
        if self._valid_rx_state():
            self._peer_interested = True
            self._client.peer_interested(self)

    def rx_not_interested(self):
        if self._valid_rx_state():
            self._peer_interested = False
            self._client.peer_not_interested(self)

    def rx_have(self, index):
        if self._valid_rx_state():
            self._client.peer_has(self, index)

    def rx_request(self, index, begin, length):
        if self._valid_rx_state():
            self._client.peer_requests(self, index, begin, length)

    def rx_piece(self, index, begin, buf):
        if self._valid_rx_state():
            self._client.peer_sent_block(self, index, begin, buf)

    def rx_cancel(self, index, begin, length):
        if self._valid_rx_state():
            self._client.peer_canceled(self, index, begin, length)

    # Client calls

    def drop_connection(self):
        self._drop_connection(False)

    def choke(self):
        if self._valid_tx_state():
            self._choked = True
            self._translator.tx_choke()

    def unchoke(self):
        if self._valid_tx_state():
            self._choked = False
            self._translator.tx_unchoke()

    def interested(self):
        if self._valid_tx_state():
            self._interested = True
            self._translator.tx_interested()

    def not_interested(self):
        if self._valid_tx_state():
            self._interested = False
            self._translator.tx_not_interested()

    def have(self, index):
        if self._valid_tx_state():
            self._translator.tx_have(index)

    def request(self, index, begin, length):
        if self._valid_tx_state():
            self._translator.tx_request(index, begin, length)

    def piece(self, index, begin, buf, offset):
        if self._valid_tx_state():
            self._translator.tx_piece(index, begin, buf)

    def cancel(self, index, begin, length):
        if self._valid_tx_state():
            self._translator.tx_cancel(index, begin, length)