def test_bytes(self):
self.assertEqual(_bencode(b''), b'0:')
self.assertEqual(_bencode(b'spam'), b'4:spam')
self.assertEqual(_bencode(b'4:spam'), b'6:4:spam')
self.assertEqual(_bencode(bytearray(b'spam')), b'4:spam')
self.assertEqual(bdecode(b'0:', True), b'')
self.assertEqual(bdecode(b'4:spam', True), b'spam')
self.assertEqual(bdecode(b'6:4:spam', True), b'4:spam')
def _decode_datagram(datagram: bytes):
msg_types = {
REQUEST_TYPE: RequestDatagram,
RESPONSE_TYPE: ResponseDatagram,
ERROR_TYPE: ErrorDatagram
}
primitive: typing.Dict = bdecode(datagram)
converted = {
str(k).encode() if not isinstance(k, bytes) else k: v
for k, v in primitive.items()
}
if converted[b'0'] in [REQUEST_TYPE, ERROR_TYPE, RESPONSE_TYPE]: # pylint: disable=unsubscriptable-object
datagram_type = converted[b'0'] # pylint: disable=unsubscriptable-object
else:
raise ValueError("invalid datagram type")
datagram_class = msg_types[datagram_type]
decoded = {
k: converted[str(i).encode()] # pylint: disable=unsubscriptable-object
for i, k in enumerate(datagram_class.required_fields)
if str(i).encode() in converted # pylint: disable=unsupported-membership-test
}
for i, _ in enumerate(OPTIONAL_FIELDS):
if str(i + OPTIONAL_ARG_OFFSET).encode() in converted:
decoded[i + OPTIONAL_ARG_OFFSET] = converted[str(
i + OPTIONAL_ARG_OFFSET).encode()]
return decoded, datagram_class
def test_dict(self):
self.assertEqual(bencode({
b'foo': 42,
b'bar': b'spam'
}), b'd3:bar4:spam3:fooi42ee')
self.assertEqual(bdecode(b'd3:bar4:spam3:fooi42ee'), {
b'foo': 42,
b'bar': b'spam'
})
def test_mixed(self):
self.assertEqual(
_bencode([[b'abc', b'127.0.0.1', 1919],
[b'def', b'127.0.0.1', 1921]]),
b'll3:abc9:127.0.0.1i1919eel3:def9:127.0.0.1i1921eee')
self.assertEqual(
bdecode(b'll3:abc9:127.0.0.1i1919eel3:def9:127.0.0.1i1921eee',
True),
[[b'abc', b'127.0.0.1', 1919], [b'def', b'127.0.0.1', 1921]])
async def test_store_to_peer(self):
loop = asyncio.get_event_loop()
with dht_mocks.mock_network_loop(loop):
node_id1 = constants.generate_id()
peer1 = KademliaProtocol(loop, PeerManager(loop), node_id1,
'1.2.3.4', 4444, 3333)
peer2 = KademliaProtocol(loop, PeerManager(loop),
constants.generate_id(), '1.2.3.5', 4444,
3333)
await loop.create_datagram_endpoint(lambda: peer1,
('1.2.3.4', 4444))
await loop.create_datagram_endpoint(lambda: peer2,
('1.2.3.5', 4444))
peer = make_kademlia_peer(node_id1, '1.2.3.4', udp_port=4444)
peer2_from_peer1 = make_kademlia_peer(peer2.node_id,
peer2.external_ip,
udp_port=peer2.udp_port)
peer2_from_peer1.update_tcp_port(3333)
peer3 = make_kademlia_peer(constants.generate_id(),
'1.2.3.6',
udp_port=4444)
store_result = await peer2.store_to_peer(b'2' * 48, peer)
self.assertEqual(store_result[0], peer.node_id)
self.assertEqual(True, store_result[1])
self.assertEqual(True,
peer1.data_store.has_peers_for_blob(b'2' * 48))
self.assertEqual(False,
peer1.data_store.has_peers_for_blob(b'3' * 48))
self.assertListEqual([peer2_from_peer1],
peer1.data_store.get_storing_contacts())
peer1.data_store.completed_blobs.add(
binascii.hexlify(b'2' * 48).decode())
find_value_response = peer1.node_rpc.find_value(peer3, b'2' * 48)
self.assertEqual(len(find_value_response[b'contacts']), 0)
self.assertSetEqual(
{b'2' * 48, b'token', b'protocolVersion', b'contacts', b'p'},
set(find_value_response.keys()))
self.assertEqual(2, len(find_value_response[b'2' * 48]))
self.assertEqual(find_value_response[b'2' * 48][0],
peer2_from_peer1.compact_address_tcp())
self.assertDictEqual(bdecode(bencode(find_value_response)),
find_value_response)
find_value_page_above_pages_response = peer1.node_rpc.find_value(
peer3, b'2' * 48, page=10)
self.assertNotIn(b'2' * 48, find_value_page_above_pages_response)
peer1.stop()
peer2.stop()
peer1.disconnect()
peer2.disconnect()
def decode_datagram(
datagram: bytes
) -> typing.Union[RequestDatagram, ResponseDatagram, ErrorDatagram]:
msg_types = {
REQUEST_TYPE: RequestDatagram,
RESPONSE_TYPE: ResponseDatagram,
ERROR_TYPE: ErrorDatagram
}
primitive: typing.Dict = bdecode(datagram)
if primitive[0] in [REQUEST_TYPE, ERROR_TYPE, RESPONSE_TYPE]: # pylint: disable=unsubscriptable-object
datagram_type = primitive[0] # pylint: disable=unsubscriptable-object
else:
raise ValueError("invalid datagram type")
datagram_class = msg_types[datagram_type]
return datagram_class(
**{
k: primitive[i] # pylint: disable=unsubscriptable-object
for i, k in enumerate(datagram_class.fields) if i in primitive # pylint: disable=unsupported-membership-test
})
def test_list(self):
self.assertEqual(_bencode([b'spam', 42]), b'l4:spami42ee')
self.assertEqual(bdecode(b'l4:spami42ee', True), [b'spam', 42])
def test_integer(self):
self.assertEqual(_bencode(42), b'i42e')
self.assertEqual(bdecode(b'i42e', True), 42)
def test_fail_bad_type(self):
with self.assertRaises(DecodeError):
bdecode(b'd4le', True)