From 2df8a1d99d5cea2999dc9898a5dfa7693782199b Mon Sep 17 00:00:00 2001 From: Victor Shyba Date: Sat, 5 Mar 2022 02:42:58 -0300 Subject: [PATCH] make a helper function to announce --- lbry/torrent/tracker.py | 19 +++++++++++++++++++ tests/unit/torrent/test_tracker.py | 8 +++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/lbry/torrent/tracker.py b/lbry/torrent/tracker.py index 14fd3c136..3d548e45b 100644 --- a/lbry/torrent/tracker.py +++ b/lbry/torrent/tracker.py @@ -4,6 +4,8 @@ import asyncio import logging from collections import namedtuple +from lbry.utils import resolve_host + log = logging.getLogger(__name__) # see: http://bittorrent.org/beps/bep_0015.html and http://xbtt.sourceforge.net/udp_tracker_protocol.html ConnectRequest = namedtuple("ConnectRequest", ["connection_id", "action", "transaction_id"]) @@ -108,3 +110,20 @@ class UDPTrackerClientProtocol(asyncio.DatagramProtocol): def connection_lost(self, exc: Exception = None) -> None: self.transport = None + + +async def get_peer_list(info_hash, node_id, port, tracker_ip, tracker_port): + 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) + return reply.peers + except asyncio.CancelledError: + raise + except Exception as exc: + log.warning("Error fetching from tracker: %s", exc.args) + return [] + finally: + transport.close() \ No newline at end of file diff --git a/tests/unit/torrent/test_tracker.py b/tests/unit/torrent/test_tracker.py index 9f1ebf106..66f147f43 100644 --- a/tests/unit/torrent/test_tracker.py +++ b/tests/unit/torrent/test_tracker.py @@ -4,7 +4,7 @@ from functools import reduce from lbry.testcase import AsyncioTestCase from lbry.torrent.tracker import UDPTrackerClientProtocol, encode, decode, CompactIPv4Peer, ConnectRequest, \ - ConnectResponse, AnnounceRequest, ErrorResponse, AnnounceResponse + ConnectResponse, AnnounceRequest, ErrorResponse, AnnounceResponse, get_peer_list class UDPTrackerServerProtocol(asyncio.DatagramProtocol): # for testing. Not suitable for production @@ -54,6 +54,12 @@ class UDPTrackerClientTestCase(AsyncioTestCase): self.assertEqual(announcement.peers, [CompactIPv4Peer(int.from_bytes(bytes([127, 0, 0, 1]), "big", signed=False), 4444)]) + 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)]) + async def test_error(self): info_hash = random.getrandbits(160).to_bytes(20, "big", signed=False) peer_id = random.getrandbits(160).to_bytes(20, "big", signed=False)