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)