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()