/
handshaketranslator.py
122 lines (97 loc) · 4.63 KB
/
handshaketranslator.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
"""
The HandshakeTranslator translates a stream of bytes into higher level messages
and vice versa related to the handshake between BitTorrent peers. The bytes
are exchanged with a readerwriter and the higher level messages are exchanged
with a receiver. The creator of the HandshakeTranslator must subsequently
provide it with a readerwriter using the set_readerwriter() method. When it
no longer needs the HandshakeTranslator, it should call the
unset_readerwriter() method so that all references to the HandshakeTranslator
can be released and it can be properly garbage collected.
On the receiver side, when told to transmit a handshake message, the
HandshakeTranslator generates a set of bytes to send to the readerwriter. When
the readerwriter presents the HandshakeTranslator with a set of bytes that
represent a handshake message, the translator notifies the receiver. It also
notifies the receiver when a non-handshake message is received or the
connection is lost.
A receiver must implement three methods, rx_handshake(), rx_non_handshake and
connection_lost().
On the readerwriter side, when incoming bytes are available, the readerwriter
asks the HandshakeTranslator 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 HandshakeTranslator to get enough bytes to translate the
next part of the message expected.
A readerwriter must implement set_receiver(), unset_receiver() and tx_bytes()
Right now, the HandshakeTranslor reports the handshake after it has received
the peer_id but documentation seems to indicate it should do this right after
receiving the info_hash.
Would it make sense to make a Translator base class with common code which all
the actual translators would inherit from?
"""
import logging
import struct
logger = logging.getLogger('bt.handshaketranslator')
_BUFFER_SIZE = 48
_LENGTH_LEN = 1
_REST_LEN = 48
class HandshakeTranslator(object):
class _States(object):
Length, Protocol, Rest = range(3)
def __init__(self):
self._buffer = bytearray(_BUFFER_SIZE)
self._view = memoryview(self._buffer)
self._rx_state = self._States.Length
self._bytes_needed = _LENGTH_LEN
self._bytes_received = 0
self._receiver = None
self._readerwriter = None
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._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('B',
buffer(self._buffer[0:_LENGTH_LEN]))
self._rx_state = self._States.Protocol
self._bytes_needed = length
self._bytes_received = 0
elif self._rx_state == self._States.Protocol:
buf = buffer(self._buffer[0:self._bytes_received])
(pstr,) = struct.unpack('{}s'.format(self._bytes_received),
buf)
if pstr == 'BitTorrent protocol':
self._rx_state = self._States.Rest
self._bytes_needed = _REST_LEN
self._bytes_received = 0
else:
if self._receiver:
self._receiver.rx_non_handshake()
else:
buf = buffer(self._buffer[0:_REST_LEN])
(reserved, info_hash, peer_id) = struct.unpack('8s20s20s', buf)
if self._receiver:
self._receiver.rx_handshake(reserved, info_hash, peer_id)
self._rx_state = self._States.Length
self._bytes_needed = _LENGTH_LEN
self._bytes_received = 0
def connection_lost(self):
if self._receiver:
self._receiver.connection_lost()
def tx_handshake(self, reserved, info_hash, peer_id):
if self._readerwriter:
bp = list('BitTorrent protocol')
self._readerwriter.tx_bytes(struct.pack('B19c', 19, *bp))
self._readerwriter.tx_bytes(struct.pack('8B', *([0]*8)))
self._readerwriter.tx_bytes(info_hash)
self._readerwriter.tx_bytes(peer_id)