async def test_nested_timeouts_concurrent(self):
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.002):
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.1):
# Pretend we crunch some numbers.
time.sleep(0.01)
await asyncio.sleep(1)
async def test_nested_timeout_in_finally(self):
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.01):
try:
await asyncio.sleep(1)
finally:
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.01):
await asyncio.sleep(10)
async def stream_torrent(loop, torrent, stream_func, filter_func):
"""
Stream torrent.
This will launch:
- A logger to log torrent alerts (while torrent is downloading)
- Torrent download
- Once torrent.started, will try to find a playable file, and
fill a ``PERCENT_CACHE`` cache to start playing
"""
async def alert_watcher(torrent):
""" Watch and log torrent alerts from LT """
while not torrent.finished:
alert = torrent.session.pop_alert()
if alert:
LOG.info(alert)
await asyncio.sleep(1)
# Force sequential mode
torrent.sequential(True)
# Start alert watcher task
loop.create_task(alert_watcher(torrent))
# If we're actually re-asking for a torrent, just relaunch streaming
if torrent.finished:
return [
wait_for_completion(torrent),
stream_func(loop, filter_func(torrent.files))
]
# Otherwise, wait for torrent to start or finish in TIMEOUT_START+
# seconds
with asyncio.timeout(TIMEOUT_START):
while not torrent.started and not torrent.finished:
await asyncio.sleep(5)
# Only download filtered file
playable_tfile = torrent.download_only(await filter_func(torrent.files))
# If filter matches nothing, stop
if not playable_tfile:
raise Exception("Could not find a playable source that matches"
" filter function")
# Try to fill the cache, if it's not filled within the first
# ``TIMEOUT_CACHE_FILL`` seconds, exit
with asyncio.timeout(TIMEOUT_CACHE_FILL):
while playable_tfile.completed_percent < PERCENT_CACHE:
await asyncio.sleep(5)
return await asyncio.gather(
[wait_for_completion(torrent),
stream_func(loop, playable_tfile)])
async def stream_torrent(loop, torrent, stream_func, filter_func):
"""
Stream torrent.
This will launch:
- A logger to log torrent alerts (while torrent is downloading)
- Torrent download
- Once torrent.started, will try to find a playable file, and
fill a ``PERCENT_CACHE`` cache to start playing
"""
async def alert_watcher(torrent):
""" Watch and log torrent alerts from LT """
while not torrent.finished:
alert = torrent.session.pop_alert()
if alert:
LOG.info(alert)
await asyncio.sleep(1)
# Force sequential mode
torrent.sequential(True)
# Start alert watcher task
loop.create_task(alert_watcher(torrent))
# If we're actually re-asking for a torrent, just relaunch streaming
if torrent.finished:
return [wait_for_completion(torrent),
stream_func(loop, filter_func(torrent.files))]
# Otherwise, wait for torrent to start or finish in TIMEOUT_START+
# seconds
with asyncio.timeout(TIMEOUT_START):
while not torrent.started and not torrent.finished:
await asyncio.sleep(5)
# Only download filtered file
playable_tfile = torrent.download_only(await filter_func(torrent.files))
# If filter matches nothing, stop
if not playable_tfile:
raise Exception("Could not find a playable source that matches"
" filter function")
# Try to fill the cache, if it's not filled within the first
# ``TIMEOUT_CACHE_FILL`` seconds, exit
with asyncio.timeout(TIMEOUT_CACHE_FILL):
while playable_tfile.completed_percent < PERCENT_CACHE:
await asyncio.sleep(5)
return await asyncio.gather([wait_for_completion(torrent),
stream_func(loop, playable_tfile)])
async def trakt_auth_device(self, advise_callback):
"""
returns the full authorization data,
as a dict, from wich we'll get "access_token"
"""
with aiohttp.ClientSession() as session:
async with session.post(
self.code_url, data={"client_id": self.cid}
) as resp:
assert resp.status == 200
auth_data = await resp.json()
await advise_callback(auth_data)
with suppress(TimeoutError):
data = {"client_id": self.cid, "client_secret": self.csecret,
"code": auth_data['device_code']}
with asyncio.timeout(auth_data['expires_in']):
with aiohttp.ClientSession() as session:
while True:
async with session.post(self.token_url, data=data) as resp:
if resp.status == 200:
return await resp.json()
else:
await asyncio.sleep(auth_data["interval"])
return False
async def trakt_auth_device(self, advise_callback):
"""
returns the full authorization data,
as a dict, from wich we'll get "access_token"
"""
with aiohttp.ClientSession() as session:
async with session.post(self.code_url,
data={"client_id": self.cid}) as resp:
assert resp.status == 200
auth_data = await resp.json()
await advise_callback(auth_data)
with suppress(TimeoutError):
data = {
"client_id": self.cid,
"client_secret": self.csecret,
"code": auth_data['device_code']
}
with asyncio.timeout(auth_data['expires_in']):
with aiohttp.ClientSession() as session:
while True:
async with session.post(self.token_url,
data=data) as resp:
if resp.status == 200:
return await resp.json()
else:
await asyncio.sleep(auth_data["interval"])
return False
def dispatch(self, event, time=10):
for action in (a for a, f in self.bindings if f.check(event)):
LOG.debug("Dispatching {}".format(str(action)))
try:
self.queue.put_nowait(
(functools.partial(action, event), timeout(time)))
except QueueFull:
LOG.error("The unbounded queue is full! Pretty weird, eh?")
async def test_nested_timeouts_loop_busy(self):
# After the inner timeout is an expensive operation which should
# be stopped by the outer timeout.
loop = asyncio.get_running_loop()
# Disable a message about long running task
loop.slow_callback_duration = 10
t0 = loop.time()
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.1): # (1)
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.01): # (2)
# Pretend the loop is busy for a while.
time.sleep(0.1)
await asyncio.sleep(1)
# TimeoutError was cought by (2)
await asyncio.sleep(
10) # This sleep should be interrupted by (1)
t1 = loop.time()
self.assertTrue(t0 <= t1 <= t0 + 1)
async def test_foreign_exception_on_timeout(self):
async def crash():
try:
await asyncio.sleep(1)
finally:
1 / 0
with self.assertRaises(ZeroDivisionError):
async with asyncio.timeout(0.01):
await crash()
async def test_timeout_zero(self):
loop = asyncio.get_running_loop()
t0 = loop.time()
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0) as cm:
await asyncio.sleep(10)
t1 = loop.time()
self.assertTrue(cm.expired())
# 2 sec for slow CI boxes
self.assertLess(t1 - t0, 2)
self.assertTrue(t0 <= cm.when() <= t1)
async def test_timeout_disabled(self):
loop = asyncio.get_running_loop()
t0 = loop.time()
async with asyncio.timeout(None) as cm:
await asyncio.sleep(0.01)
t1 = loop.time()
self.assertFalse(cm.expired())
self.assertIsNone(cm.when())
# 2 sec for slow CI boxes
self.assertLess(t1 - t0, 2)
async def test_timeout_not_called(self):
loop = asyncio.get_running_loop()
t0 = loop.time()
async with asyncio.timeout(10) as cm:
await asyncio.sleep(0.01)
t1 = loop.time()
self.assertFalse(cm.expired())
# 2 sec for slow CI boxes
self.assertLess(t1 - t0, 2)
self.assertGreater(cm.when(), t1)
async def test_waiter_cancelled(self):
loop = asyncio.get_running_loop()
cancelled = False
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.01):
try:
await asyncio.sleep(10)
except asyncio.CancelledError:
cancelled = True
raise
self.assertTrue(cancelled)
async def _stream_torrent(loop, torrent, stream_func, filter_func):
loop.create_task(alert_watcher(torrent))
if torrent.finished:
#: Parallel launch stream function and wait for completion from now on
#: We'll forget the torrent itself and relie
#: In case we already got it
return [
wait_for_completion(torrent),
stream_func(loop, filter_func(torrent.files))
]
with asyncio.timeout(10 * 60): #: TODO Make this timeout configurable
while not torrent.started and not torrent.finished:
await asyncio.sleep(5)
#: Secuential download.
torrent.handle.set_sequential_download(True)
#: Filter function must make sure to be precise...
#: It gets a list of torrent.file
playable_tfile = torrent.download_only(await filter_func(torrent.files))
LOG.info("Found playable file: %s", playable_tfile)
if not playable_tfile:
raise Exception("Could not find a playable source that matches"
" filter function")
try:
with asyncio.timeout(5 * 60):
while True:
if playable_tfile.completed_percent >= 5:
break
await asyncio.sleep(5)
except asyncio.TimeoutError:
raise Exception('Could not get playable source in time')
#: Parallel launch stream function and wait for completion from now on
#: We'll forget the torrent itself and relie
return [wait_for_completion(torrent), stream_func(loop, playable_tfile)]
async def _stream_torrent(loop, torrent, stream_func, filter_func):
loop.create_task(alert_watcher(torrent))
if torrent.finished:
#: Parallel launch stream function and wait for completion from now on
#: We'll forget the torrent itself and relie
#: In case we already got it
return [wait_for_completion(torrent),
stream_func(loop, filter_func(torrent.files))]
with asyncio.timeout(10 * 60): #: TODO Make this timeout configurable
while not torrent.started and not torrent.finished:
await asyncio.sleep(5)
#: Secuential download.
torrent.handle.set_sequential_download(True)
#: Filter function must make sure to be precise...
#: It gets a list of torrent.file
playable_tfile = torrent.download_only(await filter_func(torrent.files))
LOG.info("Found playable file: %s", playable_tfile)
if not playable_tfile:
raise Exception("Could not find a playable source that matches"
" filter function")
try:
with asyncio.timeout(5 * 60):
while True:
if playable_tfile.completed_percent >= 5:
break
await asyncio.sleep(5)
except asyncio.TimeoutError:
raise Exception('Could not get playable source in time')
#: Parallel launch stream function and wait for completion from now on
#: We'll forget the torrent itself and relie
return [wait_for_completion(torrent),
stream_func(loop, playable_tfile)]
async def client_handler(websocket, path):
# print(websocket)
async def send_string(msg):
try:
await websocket.send(msg)
except Exception as e:
print(e)
return False
return True
# Client should provide a proof of work
answ = random.randint(0, 2**ANSWER_SIZE - 1)
answ = hex(answ)[2:]
answ = answ.zfill(ANSWER_SIZE // 4)
chal = hashlib.md5(answ.encode("utf8")).hexdigest()
chal = bin(int(chal, 16))[2:]
chal = chal[:CHALLENGE_SIZE].zfill(CHALLENGE_SIZE)
if not await send_string(CHALLENGE + chal): return
try:
with asyncio.timeout(TIMEOUT):
answer = await websocket.recv()
if len(answer) != ANSWER_SIZE // 4:
print("Invalid response length.")
# await send_string(UNACCEPTABLE)
return
_chal = hashlib.md5(answer.encode("utf8")).hexdigest()
_chal = bin(int(_chal, 16))[2:]
_chal = _chal[:CHALLENGE_SIZE].zfill(CHALLENGE_SIZE)
if _chal != chal:
print("Invalid response to challenge.")
# await send_string(UNACCEPTABLE)
return
except Exception as e:
print(e)
return
# Initial permutation, persists during a session
phi = np.random.permutation(N)
i = 0
while i < N:
found = False
# Current permutation and its inverse
# Persists until a right guess, or until N/2 wrong guesses
rho = np.random.permutation(N)
rho_inv = np.argsort(rho)
result = rho[phi[rho_inv]]
if not await send_string(RE_ARRANGED + str(i).zfill(ALIGN)): return
for j in range(0, MAX_GUESSES):
try:
if not await send_string(GIVE_GUESS + str(j).zfill(ALIGN)):
return
with asyncio.timeout(TIMEOUT):
answer = await websocket.recv()
if len(answer) > 3:
raise Exception("Input too long")
guess = int(answer)
actual = result[guess]
except Exception as e:
print(e)
return
if actual == i:
if not await send_string(GREATE_GUESS + PAD): return
found = True
# print('---> GREAT!')
break
# print('--->' + str(actual))
if not await send_string(WRONG_GUESS + str(actual).zfill(ALIGN)):
return
if found:
i += 1
else:
await send_string(BYE + PAD)
return
await send_string(FLAG_IS + 'SharifCTF{flagfalgflag}')
async def test_foreign_cancel_doesnt_timeout_if_not_expired(self):
with self.assertRaises(asyncio.CancelledError):
async with asyncio.timeout(10) as cm:
asyncio.current_task().cancel()
await asyncio.sleep(10)
self.assertFalse(cm.expired())
async def test_repr_finished(self):
async with asyncio.timeout(10) as cm:
await asyncio.sleep(0)
self.assertEqual(repr(cm), "<Timeout [finished]>")
async def test_foreign_exception_passed(self):
with self.assertRaises(KeyError):
async with asyncio.timeout(0.01) as cm:
raise KeyError
self.assertFalse(cm.expired())
async def outer() -> None:
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.001):
await asyncio.sleep(10)
async def test_timeout_basic(self):
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.01) as cm:
await asyncio.sleep(10)
self.assertTrue(cm.expired())
async def test_repr_active(self):
async with asyncio.timeout(10) as cm:
self.assertRegex(repr(cm), r"<Timeout \[active\] when=\d+\.\d*>")
async def test_repr_disabled(self):
async with asyncio.timeout(None) as cm:
self.assertEqual(repr(cm), r"<Timeout [active] when=None>")
async def client_handler(websocket, path):
# print(websocket)
async def send_string(msg):
try:
await websocket.send(msg)
except Exception as e:
print(e)
return False
return True
# Client should provide a proof of work
answ = random.randint(0, 2**ANSWER_SIZE - 1)
answ = hex(answ)[2:]
answ = answ.zfill(ANSWER_SIZE//4)
chal = hashlib.md5(answ.encode("utf8")).hexdigest()
chal = bin(int(chal, 16))[2:]
chal = chal[:CHALLENGE_SIZE].zfill(CHALLENGE_SIZE)
if not await send_string(CHALLENGE + chal): return
try:
with asyncio.timeout(TIMEOUT):
answer = await websocket.recv()
if len(answer) != ANSWER_SIZE//4:
print("Invalid response length.")
# await send_string(UNACCEPTABLE)
return
_chal = hashlib.md5(answer.encode("utf8")).hexdigest()
_chal = bin(int(_chal, 16))[2:]
_chal = _chal[:CHALLENGE_SIZE].zfill(CHALLENGE_SIZE)
if _chal != chal:
print("Invalid response to challenge.")
# await send_string(UNACCEPTABLE)
return
except Exception as e:
print(e)
return
# Initial permutation, persists during a session
phi = np.random.permutation(N)
i = 0
while i < N:
found = False
# Current permutation and its inverse
# Persists until a right guess, or until N/2 wrong guesses
rho = np.random.permutation(N)
rho_inv = np.argsort(rho)
result = rho[phi[rho_inv]]
if not await send_string(RE_ARRANGED + str(i).zfill(ALIGN)): return
for j in range(0, MAX_GUESSES):
try:
if not await send_string(GIVE_GUESS + str(j).zfill(ALIGN)): return
with asyncio.timeout(TIMEOUT):
answer = await websocket.recv()
if len(answer) > 3:
raise Exception("Input too long")
guess = int(answer)
actual = result[guess]
except Exception as e:
print(e)
return
if actual == i:
if not await send_string(GREATE_GUESS + PAD): return
found = True
# print('---> GREAT!')
break
# print('--->' + str(actual))
if not await send_string(WRONG_GUESS + str(actual).zfill(ALIGN)): return
if found:
i += 1
else:
await send_string(BYE + PAD)
return
await send_string(FLAG_IS + 'SharifCTF{flagfalgflag}')
async def test_repr_expired(self):
with self.assertRaises(TimeoutError):
async with asyncio.timeout(0.01) as cm:
await asyncio.sleep(10)
self.assertEqual(repr(cm), "<Timeout [expired]>")