add arg to announce stopped, removing the announcement

This commit is contained in:
Victor Shyba 2022-03-05 02:55:19 -03:00
parent 2df8a1d99d
commit dc6f8c4fc4
2 changed files with 12 additions and 6 deletions

View file

@ -77,14 +77,15 @@ class UDPTrackerClientProtocol(asyncio.DatagramProtocol):
return decode(ConnectResponse,
await self.request(ConnectRequest(0x41727101980, 0, transaction_id), tracker_ip, tracker_port))
async def announce(self, info_hash, peer_id, port, tracker_ip, tracker_port, connection_id=None):
async def announce(self, info_hash, peer_id, port, tracker_ip, tracker_port, connection_id=None, stopped=False):
if not connection_id:
reply = await self.connect(tracker_ip, tracker_port)
connection_id = reply.connection_id
# this should make the key deterministic but unique per info hash + peer id
key = int.from_bytes(info_hash[:4], "big") ^ int.from_bytes(peer_id[:4], "big") ^ port
transaction_id = random.getrandbits(32)
req = AnnounceRequest(connection_id, 1, transaction_id, info_hash, peer_id, 0, 0, 0, 1, 0, key, -1, port)
req = AnnounceRequest(
connection_id, 1, transaction_id, info_hash, peer_id, 0, 0, 0, 3 if stopped else 1, 0, key, -1, port)
reply = await self.request(req, tracker_ip, tracker_port)
return decode(AnnounceResponse, reply), connection_id
@ -112,13 +113,13 @@ class UDPTrackerClientProtocol(asyncio.DatagramProtocol):
self.transport = None
async def get_peer_list(info_hash, node_id, port, tracker_ip, tracker_port):
async def get_peer_list(info_hash, node_id, port, tracker_ip, tracker_port, stopped=False):
node_id = node_id or random.getrandbits(160).to_bytes(20, "big", signed=False)
tracker_ip = await resolve_host(tracker_ip, tracker_port, 'udp')
proto = UDPTrackerClientProtocol()
transport, _ = await asyncio.get_running_loop().create_datagram_endpoint(lambda: proto, local_addr=("0.0.0.0", 0))
try:
reply, _ = await proto.announce(info_hash, node_id, port, tracker_ip, tracker_port)
reply, _ = await proto.announce(info_hash, node_id, port, tracker_ip, tracker_port, stopped=stopped)
return reply.peers
except asyncio.CancelledError:
raise

View file

@ -32,7 +32,11 @@ class UDPTrackerServerProtocol(asyncio.DatagramProtocol): # for testing. Not su
else:
self.peers.setdefault(req.info_hash, [])
compact_ip = reduce(lambda buff, x: buff + bytearray([int(x)]), address[0].split('.'), bytearray())
self.peers[req.info_hash].append(compact_ip + req.port.to_bytes(2, "big", signed=False))
compact_address = compact_ip + req.port.to_bytes(2, "big", signed=False)
if req.event != 3:
self.peers[req.info_hash].append(compact_address)
elif compact_address in self.peers[req.info_hash]:
self.peers[req.info_hash].remove(compact_address)
peers = [decode(CompactIPv4Peer, peer) for peer in self.peers[req.info_hash]]
resp = encode(AnnounceResponse(1, req.transaction_id, 1700, 0, len(peers), peers))
return self.transport.sendto(resp, address)
@ -57,8 +61,9 @@ class UDPTrackerClientTestCase(AsyncioTestCase):
async def test_announce_using_helper_function(self):
info_hash = random.getrandbits(160).to_bytes(20, "big", signed=False)
peers = await get_peer_list(info_hash, None, 4444, "127.0.0.1", 59900)
self.assertEqual(len(peers), 1)
self.assertEqual(peers, [CompactIPv4Peer(int.from_bytes(bytes([127, 0, 0, 1]), "big", signed=False), 4444)])
peers = await get_peer_list(info_hash, None, 4444, "127.0.0.1", 59900, stopped=True)
self.assertEqual(peers, [])
async def test_error(self):
info_hash = random.getrandbits(160).to_bytes(20, "big", signed=False)