def test_bencode_bdecode(self): # Decode whole Torrent files torrent_dir = os.path.join(DATA_DIR, "torrent_files") for file in os.listdir(torrent_dir): filename = os.path.join(torrent_dir, os.fsdecode(file)) with open(filename, "rb") as f: with self.subTest(filename=filename): s = f.read() self.assertEqual(s, bencode(bdecode(s)))
def test_Bencode_failure(self): fail_cases = [ # Number but not an integer 1.0, # String instead of a bytestring "a", # Dictionary keys must be bytestrings { 1: b"a", 2: b"b" }, { "b": b"a", "a": b"b" } ] for testCase in fail_cases: with self.subTest(testCase=testCase): with self.assertRaises(ValueError, msg="case '{0}'".format(testCase)): bencode(testCase)
def _from_file(cls, torrent_file: str): """Create an instance of Torrent from a metainfo file""" with open(torrent_file, "rb") as f: m = bdecode(f.read()) # TODO: check type of optional items if validate_structure(m, METAINFO_SINGLE_FILE): single_file_mode = True elif validate_structure(m, METAINFO_MULTIPLE_FILES): single_file_mode = False else: raise ValueError("Invalid torrent file") info = m[b'info'] info_hash = sha1(bencode(info)).digest() if b'announce-list' in m: announce = m[b'announce-list'] else: announce = [[m[b'announce']]] announce_decoded = [] for trackers in announce: announce_decoded.append(list(map(lambda x: x.decode('utf-8'), trackers))) name = info[b'name'].decode("utf-8") hashes = split(info[b'pieces'], sha1().digest_size) piece_length = info[b'piece length'] contents = [] if single_file_mode: contents.append(("", sanitize(name), info[b"length"])) else: for file_dict in info[b'files']: path_str = [sanitize(b.decode("utf-8")) for b in file_dict[b'path']] directory = os.path.join("", *path_str[:-1]) filename = path_str[-1] contents.append((directory, filename, file_dict[b'length'])) return cls(announce_decoded, name, contents, info_hash, hashes, piece_length)
def two_peer_swarm(data_dir, seeder_port=6881, leecher_port=6882): """ - Create a torrent file for a given directory - Create two Torrent instances from the file - Exchange data between these two instances until the leecher has downloaded the whole file """ # Setup asyncio loop loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.set_debug(False) executor = concurrent.futures.ThreadPoolExecutor() loop.set_default_executor(executor) # Setup working directories tmp_dir = tempfile.mkdtemp() seeder_dir = os.path.join(tmp_dir, "seeder") leecher_dir = os.path.join(tmp_dir, "leecher") # The seeder's directory is initiated with all the files shutil.copytree(data_dir, seeder_dir) # The leecher's directory is empty os.makedirs(leecher_dir, exist_ok=True) # Create metainfo file on disk on the fly torrent_file = os.path.join(tmp_dir, 'metainfo.torrent') m = metainfo(data_dir, 32768, [['http://www.example.com/announce']]) with open(torrent_file, 'wb') as f: f.write(bencode(m)) async def mock_get_peers(r, *_): return r # Start seeder and override get_peers method to return an empty list torrent_seeder = loop.run_until_complete( Torrent.create(torrent_file, seeder_dir)) torrent_seeder.tracker.get_peers = functools.partial(mock_get_peers, []) # Start leecher and override get_peers method to return the address of the seeder torrent_leecher = loop.run_until_complete( Torrent.create(torrent_file, leecher_dir)) torrent_leecher.tracker.get_peers = functools.partial( mock_get_peers, [("127.0.0.1", seeder_port)]) async def wait_for(torrent: Torrent, events: List[str]): event = None while event not in events: event = await torrent.queue.get() return event # Futures f_seeder = asyncio.ensure_future(torrent_seeder.download(seeder_port)) f_wait_accept_conns = asyncio.ensure_future( wait_for(torrent_seeder, ["EVENT_ACCEPT_CONNECTIONS", "EVENT_END"])) f_download_complete = None futures = {f_seeder, f_wait_accept_conns} while futures: done, futures = loop.run_until_complete( asyncio.wait(futures, return_when=asyncio.FIRST_COMPLETED)) for item in done: result = item.result() # Once the seeder accepts connection, start the leecher and wait for it to complete # the download if item == f_wait_accept_conns: if result == "EVENT_ACCEPT_CONNECTIONS": f_leecher = asyncio.ensure_future( torrent_leecher.download(leecher_port), loop=loop) f_download_complete = asyncio.ensure_future( wait_for(torrent_leecher, ["EVENT_DOWNLOAD_COMPLETE", "EVENT_END"])) futures = futures | {f_leecher, f_download_complete} elif result == "EVENT_END": print("leecher failed") # Once the leecher has downloaded the file, stop all torrents if item == f_download_complete: if result == "EVENT_DOWNLOAD_COMPLETE": loop.run_until_complete(torrent_leecher.stop()) loop.run_until_complete(torrent_seeder.stop()) elif result == "EVENT_END": print("seeder failed") loop.stop() loop.close() assert filecmp.dircmp(seeder_dir, data_dir) assert filecmp.dircmp(leecher_dir, data_dir) shutil.rmtree(tmp_dir)
def test_Bencode_success(self): for (key, value) in validBEncodings.items(): with self.subTest(case=value, expected=key, result=bencode(value)): self.assertEqual(bencode(value), key)