add arg to announce stopped, removing the announcement
This commit is contained in:
parent
2df8a1d99d
commit
dc6f8c4fc4
2 changed files with 12 additions and 6 deletions
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in a new issue