def test_insert_secret_with_system_multicall(self):
# create client with secret
secret = "hello"
client = Client(secret=secret)
responses.add_callback(responses.POST,
client.server,
callback=self.call_params_callback)
# create params
params = [[
{
"methodName": client.ADD_URI,
"params": ["param1", "param2"]
},
{
"methodName": client.ADD_URI,
"params": ["param3", "param4"]
},
]]
# copy params and insert secret
expected_params = deepcopy(params)
for param in expected_params[0]:
param["params"].insert(0, f"token:{secret}")
# call function and assert result
resp = client.call(client.MULTICALL, params, insert_secret=True)
assert resp == expected_params
def test_insert_secret_with_multicall2(self):
# create client with secret
secret = "hello"
client = Client(secret=secret)
responses.add_callback(responses.POST,
client.server,
callback=self.call_params_callback)
# create params
params_1 = ["2089b05ecca3d829"]
params_2 = ["2fa07b6e85c40205"]
calls = [(client.REMOVE, params_1), (client.REMOVE, params_2)]
# copy params and insert secret
expected_params = [[
{
"methodName": client.REMOVE,
"params": deepcopy(params_1)
},
{
"methodName": client.REMOVE,
"params": deepcopy(params_2)
},
]]
for param in expected_params[0]:
param["params"].insert(0, f"token:{secret}")
# call function and assert result
resp = client.multicall2(calls, insert_secret=True)
assert resp == expected_params
def test_multicall2(self):
client = Client()
responses.add_callback(responses.POST,
client.server,
callback=self.call_params_callback)
# create params
params_1 = ["2089b05ecca3d829"]
params_2 = ["2fa07b6e85c40205"]
calls = [(client.REMOVE, params_1), (client.REMOVE, params_2)]
# copy params and insert secret
expected_params = [[
{
"methodName": client.REMOVE,
"params": deepcopy(params_1)
},
{
"methodName": client.REMOVE,
"params": deepcopy(params_2)
},
]]
# call function and assert result
resp = client.multicall2(calls)
assert resp == expected_params
def test_call_raises_custom_error(self):
client = Client()
responses.add(
responses.POST, client.server, json={"error": {"code": 1, "message": "Custom message"}}, status=200
)
with pytest.raises(ClientException, match=r"Custom message") as e:
client.call("aria2.method")
assert e.code == 1
def test_call_raises_known_error(self):
client = Client()
responses.add(
responses.POST,
client.server,
json={"error": {"code": JSONRPC_PARSER_ERROR, "message": "Custom message"}},
status=200,
)
with pytest.raises(ClientException, match=rf"{JSONRPC_CODES[JSONRPC_PARSER_ERROR]}\nCustom message") as e:
client.call("aria2.method")
assert e.code == JSONRPC_PARSER_ERROR
def test_does_not_insert_secret_if_told_so(self):
# create client with secret
secret = "hello"
client = Client(secret=secret)
responses.add_callback(responses.POST, client.server, callback=self.call_params_callback)
# create params
params = ["param1", "param2"]
# call function and assert result
resp = client.call("other.method", params, insert_secret=False)
assert secret not in resp
def test_batch_call(self):
client = Client()
responses.add_callback(responses.POST, client.server, callback=self.batch_call_params_callback)
# create params
params_1 = ["param1", "param2"]
params_2 = ["param3", "param4"]
# copy params and insert secret
expected_params = [params_1, params_2]
# call function and assert result
resp = client.batch_call([(client.ADD_URI, params_1, 0), (client.ADD_METALINK, params_2, 1)])
assert resp == expected_params
def test_insert_secret_with_aria2_method_call(self):
# create client with secret
secret = "hello"
client = Client(secret=secret)
responses.add_callback(responses.POST, client.server, callback=self.call_params_callback)
# create params
params = ["param1", "param2"]
# copy params and insert secret
expected_params = deepcopy(params)
expected_params.insert(0, secret)
# call function and assert result
resp = client.call(client.ADD_URI, params, insert_secret=True)
assert resp == expected_params
def test_insert_secret_with_batch_call(self):
# create client with secret
secret = "hello"
client = Client(secret=secret)
responses.add_callback(responses.POST, client.server, callback=self.batch_call_params_callback)
# create params
params_1 = ["param1", "param2"]
params_2 = ["param3", "param4"]
# copy params and insert secret
expected_params = [deepcopy(params_1), deepcopy(params_2)]
for p in expected_params:
p.insert(0, secret)
# call function and assert result
resp = client.batch_call(
[(client.ADD_URI, params_1, 0), (client.ADD_METALINK, params_2, 1)], insert_secret=True
)
assert resp == expected_params
def __init__(self, tmp_dir, port, config=None, session=None, secret=""):
self.tmp_dir = tmp_dir
self.port = port
# create the command used to launch an aria2c process
command = [
"aria2c",
f"--dir={self.tmp_dir}",
"--file-allocation=none",
"--quiet",
"--enable-rpc=true",
f"--rpc-listen-port={self.port}",
]
if config:
command.append(f"--conf-path={config}")
else:
# command.append("--no-conf")
config = CONFIGS_DIR / "default.conf"
command.append(f"--conf-path={config}")
if session:
if isinstance(session, list):
session_path = self.tmp_dir / "_session.txt"
with open(session_path, "w") as stream:
stream.write("\n".join(session))
command.append(f"--input-file={session_path}")
else:
session_path = SESSIONS_DIR / session
if not session_path.exists():
raise ValueError(f"no such session: {session}")
command.append(f"--input-file={session_path}")
if secret:
command.append(f"--rpc-secret={secret}")
self.command = command
self.process = None
# create the client with port
self.client = Client(port=self.port, secret=secret, timeout=20)
# create the API instance
self.api = API(self.client)
def __init__(self,
host=DEFAULT_HOST,
port=DEFAULT_PORT,
secret="",
poll=1.0,
global_options=None): # nosec
# type: (str, int, str, float, Optional[dict]) -> None
"""Initialize the aria2 client with `host`, `port` and `secret`.
Poll aria2 every `poll` seconds.
"""
self.aria2 = Client(host, port, secret)
self.poll = poll
if global_options:
global_options = self.default_global_options.copy()
global_options.update(global_options)
self.aria2.change_global_option(global_options)
else:
self.aria2.change_global_option(self.default_global_options)
self.gids = set() # type: Set[str]
def test_listen_to_notifications_no_server(self):
client = Client(port=7035)
client.listen_to_notifications(timeout=1)
def test_client_str_returns_client_server(self):
host = "https://example.com/"
port = 7100
client = Client(host, port)
assert client.server == f"{host.rstrip('/')}:{port}/jsonrpc" == str(
client)
def test_listen_to_notifications_then_stop(port):
api = API(Client(port=port))
api.listen_to_notifications(threaded=True, timeout=1)
api.stop_listening()
assert api.listener is None
class AriaDownloader:
default_global_options = {
"max-concurrent-downloads": 5,
"remote-time": True,
}
default_options = {
"max-connection-per-server": 1,
"split": 5,
"always-resume": True,
}
max_num_results = 100
""" Download manager which uses aria2 instance to actually download the files.
Tries to respect other users of the same instance and doesn't interfere with them.
"""
def __init__(self,
host=DEFAULT_HOST,
port=DEFAULT_PORT,
secret="",
poll=1.0,
global_options=None): # nosec
# type: (str, int, str, float, Optional[dict]) -> None
"""Initialize the aria2 client with `host`, `port` and `secret`.
Poll aria2 every `poll` seconds.
"""
self.aria2 = Client(host, port, secret)
self.poll = poll
if global_options:
global_options = self.default_global_options.copy()
global_options.update(global_options)
self.aria2.change_global_option(global_options)
else:
self.aria2.change_global_option(self.default_global_options)
self.gids = set() # type: Set[str]
def query(self, method, *args, **kwargs):
try:
return getattr(self.aria2, method)(*args, **kwargs)
except ClientException as e:
if e.code == 1:
# either our code is bad, or some external actor removed our gid from aria
raise InconsistentState(e.message)
else:
raise
except ConnectionError as e:
raise ExternalProcedureUnavailable(e)
def pause_all(self):
# type: () -> None
for gid in self.gids:
self.query("pause", gid)
def resume_all(self):
# type: () -> None
for gid in self.gids:
self.query("unpause", gid)
def _entries(self, entries):
return {
entry["gid"]: entry
for entry in entries if entry["gid"] in self.gids
}
def managed_downloads(self):
# type: () -> int
return len(self.gids)
def block_one(self, progress_file=sys.stdout):
# type: (Optional[TextIO], ) -> Tuple[str, str]
"""Blocks until one download is done."""
while True:
"""fixme: This loop has a (not very serious) race condition.
For example a download might change its status from waiting to active
during the two queries. Then it would not be found.
This cannot be fixed by changing the query order as the status change graph
(waiting<->active->stopped) cannot be sorted topologically.
And even if you could, returning stopped downloads has priority over
waiting/active ones, as otherwise we would block for too long.
Does using a multi/batch-call fix this?
"""
entries = self._entries(
self.query("tell_stopped", 0,
self.max_num_results)) # complete or error
if entries:
gid, entry = entries.popitem()
try:
if entry["status"] == "complete":
assert len(entry["files"]) == 1
return gid, entry["files"][0]["path"]
elif entry["status"] == "error":
raise DownloadFailed(gid, entry["errorCode"],
entry["errorMessage"])
else:
raise RuntimeError("Unexpected status: {}".format(
entry["status"]))
finally:
self.remove_stopped(gid)
entries = self._entries(self.query("tell_active")) # active
if entries:
if progress_file:
completed = sum(
int(entry["completedLength"])
for entry in entries.values())
total = sum(
int(entry["totalLength"])
for entry in entries.values())
speed = sum(
int(entry["downloadSpeed"])
for entry in entries.values())
print(
f"{len(entries)} downloads: {completed}/{total} bytes {speed} bytes/sec",
file=progress_file,
end="\r",
)
sleep(self.poll)
continue
entries = self._entries(
self.query("tell_waiting", 0,
self.max_num_results)) # waiting or paused
if entries:
print(f"{len(entries)} downloads waiting or paused", end="\r")
sleep(self.poll)
continue
if self.gids:
"""Actually only this check is sensitive to race condition, as the looping logic
would care of retrying otherwise.
However this is the only way to check for external modifications.
"""
raise InconsistentState(
"Some downloads got lost. We either encoutered a race condition \
or some external actor removed the download")
raise WouldBlockForever("No downloads active or waiting")
def block_all(self):
# type: () -> List[Tuple[Any, Any]]
ret = [] # type: List[Tuple[Any, Any]]
while True:
try:
gid, path = self.block_one()
ret.append((None, path))
except WouldBlockForever:
break
except DownloadFailed as e:
ret.append((e.args, None))
return ret
def remove_stopped(self, gid):
# type: (str, ) -> None
self.gids.remove(gid)
# removes a complete/error/removed download
# fails on active/waiting/paused (on CTRL-C for example)
self.query("remove_download_result", gid)
def block_gid(self, gid, progress_file=sys.stdout):
# type: (str, Optional[TextIO]) -> str
"""Blocks until download is done.
If progress printing is not needed, `progress_file` should be set to `None`.
Returns the path of the downloaded file on disk.
"""
if gid not in self.gids:
raise KeyError("Invalid GID")
try:
while True:
s = self.query("tell_status", gid)
status = s["status"]
if status == "active":
if progress_file:
print(
s["completedLength"],
"/",
s["totalLength"],
"bytes",
s["downloadSpeed"],
"bytes/sec",
file=progress_file,
end="\r",
)
elif status == "waiting":
if progress_file:
print("waiting", file=progress_file, end="\r")
elif status == "paused":
if progress_file:
print("paused", file=progress_file, end="\r")
elif status == "error":
raise DownloadFailed(
gid, s["errorCode"], s["errorMessage"]
) # RuntimeError: No URI available. errorCode=8 handle
elif status == "complete":
assert len(s["files"]) == 1
return s["files"][0]["path"]
elif status == "removed":
raise RuntimeError("Someone removed our download...")
else:
raise RuntimeError(f"Unexpected status: {status}")
sleep(self.poll)
finally:
self.remove_stopped(gid)
def download(
self,
uri,
path=None,
filename=None,
headers=None,
max_connections=None,
split=None,
continue_=None,
retry_wait=None,
):
# type: (str, Optional[str], Optional[str], Optional[Sequence[str]], Optional[int], Optional[int], Optional[bool], Optional[int]) -> str
"""Downloads `uri` to directory `path`.
Does not block. Returns a download identifier.
"""
if headers is not None:
assert_type("headers", headers, Sequence)
options = self.default_options.copy()
update(
options,
{
"dir": path,
"out": filename,
"max-connection-per-server": max_connections,
"split": split,
"continue": aria_bool(continue_),
"retry-wait": retry_wait,
"header": headers,
},
)
gid = self.query("add_uri", [uri], options)
self.gids.add(gid)
return gid
def download_x(
self,
num,
uri,
path=None,
filename=None,
headers=None,
max_connections=None,
split=None,
continue_=None,
retry_wait=None,
):
# type: (int, str, Optional[str], Optional[str], Optional[Sequence[str]], Optional[int], Optional[int], Optional[bool], Optional[int]) -> Optional[Tuple[str, str, str]]
queued_gid = self.download(uri, path, filename, headers,
max_connections, split, continue_,
retry_wait)
if self.managed_downloads() >= num:
finished_gid, path = self.block_one()
return queued_gid, finished_gid, path
return None
class _Aria2Server:
def __init__(self, tmp_dir, port, config=None, session=None, secret=""):
self.tmp_dir = tmp_dir
self.port = port
# create the command used to launch an aria2c process
command = [
"aria2c",
f"--dir={self.tmp_dir}",
"--file-allocation=none",
"--quiet",
"--enable-rpc=true",
f"--rpc-listen-port={self.port}",
]
if config:
command.append(f"--conf-path={config}")
else:
# command.append("--no-conf")
config = CONFIGS_DIR / "default.conf"
command.append(f"--conf-path={config}")
if session:
if isinstance(session, list):
session_path = self.tmp_dir / "_session.txt"
with open(session_path, "w") as stream:
stream.write("\n".join(session))
command.append(f"--input-file={session_path}")
else:
session_path = SESSIONS_DIR / session
if not session_path.exists():
raise ValueError(f"no such session: {session}")
command.append(f"--input-file={session_path}")
if secret:
command.append(f"--rpc-secret={secret}")
self.command = command
self.process = None
# create the client with port
self.client = Client(port=self.port, secret=secret, timeout=20)
# create the API instance
self.api = API(self.client)
def start(self):
while True:
# create the subprocess
self.process = subprocess.Popen(self.command)
# make sure the server is running
retries = 5
while retries:
try:
self.client.list_methods()
except requests.ConnectionError:
time.sleep(0.1)
retries -= 1
else:
break
if retries:
break
def wait(self):
while True:
try:
self.process.wait()
except subprocess.TimeoutExpired:
pass
else:
break
def terminate(self):
self.process.terminate()
self.wait()
def kill(self):
self.process.kill()
self.wait()
def rmdir(self, directory=None):
if directory is None:
directory = self.tmp_dir
for item in directory.iterdir():
if item.is_dir():
self.rmdir(item)
else:
item.unlink()
directory.rmdir()
def destroy(self, force=False):
if force:
self.kill()
else:
self.terminate()
self.rmdir()
