def test_basic_built(tmp_custom_config, config_changes_dict): # noqa F811 for section in config_changes_dict["replace"]: for option in config_changes_dict["replace"][section]: if section not in configuration: configuration.add_section(section) configuration.set(section, option, config_changes_dict["replace"][section][option]) for section in config_changes_dict["remove"]: if "section" in config_changes_dict["remove"][section]: configuration.remove_section(section) else: for option in config_changes_dict["remove"][section]: configuration.remove_option(section, option) tmp_custom_config.save() if "expected_exception" in config_changes_dict: if "duplicate_exception" in config_changes_dict and config_changes_dict[ "duplicate_exception"]: with open(tmp_custom_config.config_file_path, "r") as file: content = file.read() with open(tmp_custom_config.config_file_path, "w") as file: file.write(content) file.write(content) with pytest.raises(config_changes_dict["expected_exception"]): Dispatcher(None, tmp_custom_config.config_file_path) else: Dispatcher(None, tmp_custom_config.config_file_path)
async def test_run_once(test_config: FaradayTestConfig, tmp_default_config, test_logger_handler, test_logger_folder, executor_options): # Config workspace = test_config.workspace if "workspace" not in executor_options else executor_options["workspace"] configuration.set(Sections.SERVER, "api_port", str(test_config.client.port)) configuration.set(Sections.SERVER, "host", test_config.client.host) configuration.set(Sections.SERVER, "workspace", workspace) configuration.set(Sections.TOKENS, "registration", test_config.registration_token) configuration.set(Sections.TOKENS, "agent", test_config.agent_token) path_to_basic_executor = ( Path(__file__).parent.parent / 'data' / 'basic_executor.py' ) executor_names = ["ex1"] + ([] if "extra" not in executor_options else executor_options["extra"]) configuration.set(Sections.AGENT, "executors", ",".join(executor_names)) for executor_name in executor_names: executor_section = Sections.EXECUTOR_DATA.format(executor_name) params_section = Sections.EXECUTOR_PARAMS.format(executor_name) varenvs_section = Sections.EXECUTOR_VARENVS.format(executor_name) for section in [executor_section, params_section, varenvs_section]: if section not in configuration: configuration.add_section(section) configuration.set(executor_section, "cmd", "python {}".format(path_to_basic_executor)) configuration.set(params_section, "out", "True") [configuration.set(params_section, param, "False") for param in [ "count", "spare", "spaced_before", "spaced_middle", "err", "fails"]] if "varenvs" in executor_options: for varenv in executor_options["varenvs"]: configuration.set(varenvs_section, varenv, executor_options["varenvs"][varenv]) max_size = str(64 * 1024) if "max_size" not in executor_options else executor_options["max_size"] configuration.set(executor_section, "max_size", max_size) tmp_default_config.save() async def ws_messages_checker(msg): msg_ = json.loads(msg) assert msg_ in executor_options["ws_responses"] executor_options["ws_responses"].remove(msg_) # Init and register it dispatcher = Dispatcher(test_config.client.session, tmp_default_config.config_file_path) await dispatcher.run_once(json.dumps(executor_options["data"]), ws_messages_checker) history = test_logger_handler.history assert len(executor_options["ws_responses"]) == 0 for l in executor_options["logs"]: min_count = 1 if "min_count" not in l else l["min_count"] max_count = sys.maxsize if "max_count" not in l else l["max_count"] assert max_count >= \ len(list(filter(lambda x: x.levelname == l["levelname"] and l["msg"] in x.message, history))) >= \ min_count, l["msg"]
async def test_merge_config( test_config: FaradayTestConfig, # noqa F811 tmp_default_config, # noqa F811 test_logger_handler, # noqa F811 test_logger_folder, # noqa F811 ): configuration.set(Sections.SERVER, "api_port", str(test_config.client.port)) configuration.set(Sections.SERVER, "websocket_port", str(test_config.client.port)) if test_config.base_route: configuration.set(Sections.SERVER, "base_route", test_config.base_route) configuration.set(Sections.SERVER, "ssl", str(test_config.is_ssl)) if test_config.is_ssl: configuration.set(Sections.SERVER, "ssl_cert", str(test_config.ssl_cert_path / "ok.crt")) configuration.set(Sections.SERVER, "host", "localhost") else: configuration.set(Sections.SERVER, "host", test_config.client.host) configuration.set(Sections.SERVER, "workspaces", test_config.workspaces_str()) random_workspace_name = fuzzy_string(15) configuration.set(Sections.SERVER, "workspace", random_workspace_name) test_config.workspaces = [random_workspace_name] + test_config.workspaces configuration.set(Sections.TOKENS, "registration", test_config.registration_token) configuration.set(Sections.TOKENS, "agent", test_config.agent_token) path_to_basic_executor = Path(__file__).parent.parent / "data" / "basic_executor.py" configuration.set(Sections.AGENT, "executors", "ex1,ex2,ex3,ex4") for executor_name in ["ex1", "ex3", "ex4"]: executor_section = Sections.EXECUTOR_DATA.format(executor_name) params_section = Sections.EXECUTOR_PARAMS.format(executor_name) for section in [executor_section, params_section]: if section not in configuration: configuration.add_section(section) configuration.set(executor_section, "cmd", "python {}".format(path_to_basic_executor)) configuration.set(Sections.EXECUTOR_PARAMS.format("ex1"), "param1", "True") configuration.set(Sections.EXECUTOR_PARAMS.format("ex1"), "param2", "False") configuration.set(Sections.EXECUTOR_PARAMS.format("ex3"), "param3", "False") configuration.set(Sections.EXECUTOR_PARAMS.format("ex3"), "param4", "False") tmp_default_config.save() dispatcher = Dispatcher(test_config.client.session, tmp_default_config.config_file_path) test_config.ws_data["run_data"] = {"agent_id": 1} test_config.ws_data["ws_responses"] = [{"error": "'action' key is mandatory in this websocket connection"}] test_config.executors = get_merge_executors() await dispatcher.register() await dispatcher.connect() assert len(test_config.ws_data["ws_responses"]) == 0
async def test_connect(test_config: FaradayTestConfig, tmp_default_config, test_logger_handler, test_logger_folder): configuration.set(Sections.SERVER, "api_port", str(test_config.client.port)) configuration.set(Sections.SERVER, "host", test_config.client.host) configuration.set(Sections.SERVER, "workspace", test_config.workspace) configuration.set(Sections.TOKENS, "registration", test_config.registration_token) configuration.set(Sections.TOKENS, "agent", test_config.agent_token) path_to_basic_executor = ( Path(__file__).parent.parent / 'data' / 'basic_executor.py' ) configuration.set(Sections.AGENT, "executors", "ex1,ex2,ex3") for executor_name in ["ex1","ex2","ex3"]: executor_section = Sections.EXECUTOR_DATA.format(executor_name) params_section = Sections.EXECUTOR_PARAMS.format(executor_name) for section in [executor_section, params_section]: if section not in configuration: configuration.add_section(section) configuration.set(executor_section, "cmd", "python {}".format(path_to_basic_executor)) configuration.set(Sections.EXECUTOR_PARAMS.format("ex1"), "param1", "True") configuration.set(Sections.EXECUTOR_PARAMS.format("ex1"), "param2", "False") configuration.set(Sections.EXECUTOR_PARAMS.format("ex2"), "param3", "False") configuration.set(Sections.EXECUTOR_PARAMS.format("ex2"), "param4", "False") tmp_default_config.save() dispatcher = Dispatcher(test_config.client.session, tmp_default_config.config_file_path) ws_responses = [{ 'action': 'JOIN_AGENT', 'workspace': test_config.workspace, 'token': None, 'executors': [ { "executor_name": "ex1", "args": { "param1": True, "param2": False } }, { "executor_name": "ex2", "args": { "param3": False, "param4": False } }, { "executor_name": "ex3", "args": {} } ] }] async def ws_messages_checker(msg): msg_ = json.loads(msg) assert msg_ in ws_responses ws_responses.remove(msg_) await dispatcher.connect(ws_messages_checker) assert len(ws_responses) == 0
async def test_start_and_register( register_options, test_config: FaradayTestConfig, # noqa F811 tmp_default_config, # noqa F811 test_logger_handler, # noqa F811 ): os.environ["DISPATCHER_TEST"] = "True" if "use_ssl" in register_options: if (register_options["use_ssl"] and not test_config.is_ssl) or (not register_options["use_ssl"] and test_config.is_ssl): pytest.skip( f"This test should be skipped: server_ssl:{test_config.is_ssl}" f" and config_use_ssl:f {register_options['use_ssl']}") client = test_config.client if test_config.base_route: configuration.set(Sections.SERVER, "base_route", test_config.base_route) # Config configuration.set(Sections.SERVER, "ssl", str(test_config.is_ssl)) if test_config.is_ssl: configuration.set(Sections.SERVER, "ssl_cert", str(test_config.ssl_cert_path / "ok.crt")) configuration.set(Sections.SERVER, "host", "localhost") else: configuration.set(Sections.SERVER, "host", client.host) configuration.set(Sections.SERVER, "api_port", str(client.port)) configuration.set(Sections.SERVER, "workspaces", test_config.workspaces_str()) configuration.set(Sections.EXECUTOR_DATA.format("ex1"), "cmd", "exit 1") for section in register_options["replace_data"]: for option in register_options["replace_data"][section]: if section not in configuration: configuration.add_section(section) configuration.set( section, option, register_options["replace_data"][section][option]) tmp_default_config.save() # Init and register it dispatcher = Dispatcher(client.session, tmp_default_config.config_file_path) if "expected_exception" not in register_options: await dispatcher.register(test_config.registration_token) # Control tokens assert dispatcher.agent_token == test_config.agent_token signer = TimestampSigner(test_config.app_config["SECRET_KEY"], salt="websocket_agent") agent_id = int( signer.unsign(dispatcher.websocket_token).decode("utf-8")) assert test_config.agent_id == agent_id else: if "bad_registration_token" in register_options: if register_options["bad_registration_token"] is None: token = None elif register_options["bad_registration_token"] == "incorrect": token = f"{((int(test_config.registration_token) + 1) % 1000000):06}" elif register_options["bad_registration_token"] == "bad format": token = "qewqwe" else: # == "bad" token = test_config.registration_token[0:3] else: token = test_config.registration_token with pytest.raises(register_options["expected_exception"]): await dispatcher.register(token) history = test_logger_handler.history logs_ok, failed_logs = await check_logs(history, register_options["logs"]) if "optional_logs" in register_options and not logs_ok: logs_ok, new_failed_logs = await check_logs( history, register_options["optional_logs"]) failed_logs = {"logs": failed_logs, "optional_logs": new_failed_logs} assert logs_ok, failed_logs
async def test_run_once( test_config: FaradayTestConfig, # noqa F811 tmp_default_config, # noqa F811 test_logger_handler, # noqa F811 test_logger_folder, # noqa F811 executor_options, ): # Config if "workspaces" in executor_options: test_config.workspaces = executor_options["workspaces"].split(",") workspaces = test_config.workspaces workspaces_str = test_config.workspaces_str() if test_config.base_route: configuration.set(Sections.SERVER, "base_route", test_config.base_route) configuration.set(Sections.SERVER, "api_port", str(test_config.client.port)) configuration.set(Sections.SERVER, "websocket_port", str(test_config.client.port)) configuration.set(Sections.SERVER, "workspaces", workspaces_str) if Sections.TOKENS not in configuration: configuration.add_section(Sections.TOKENS) configuration.set(Sections.TOKENS, "agent", test_config.agent_token) configuration.set(Sections.SERVER, "ssl", str(test_config.is_ssl)) if test_config.is_ssl: configuration.set(Sections.SERVER, "ssl_cert", str(test_config.ssl_cert_path / "ok.crt")) configuration.set(Sections.SERVER, "host", "localhost") else: configuration.set(Sections.SERVER, "host", test_config.client.host) path_to_basic_executor = Path( __file__).parent.parent / "data" / "basic_executor.py" executor_names = ["ex1"] + ([] if "extra" not in executor_options else executor_options["extra"]) configuration.set(Sections.AGENT, "executors", ",".join(executor_names)) test_config.executors = [] for executor_name in executor_names: executor_section = Sections.EXECUTOR_DATA.format(executor_name) params_section = Sections.EXECUTOR_PARAMS.format(executor_name) varenvs_section = Sections.EXECUTOR_VARENVS.format(executor_name) for section in [executor_section, params_section, varenvs_section]: if section not in configuration: configuration.add_section(section) configuration.set(executor_section, "cmd", "python {}".format(path_to_basic_executor)) false_params = [ "count", "spare", "spaced_before", "spaced_middle", "err", "fails", ] [ configuration.set(params_section, param, "False") for param in false_params ] configuration.set(params_section, "out", "True") if "varenvs" in executor_options: for varenv in executor_options["varenvs"]: configuration.set(varenvs_section, varenv, executor_options["varenvs"][varenv]) max_size = str( 64 * 1024) if "max_size" not in executor_options else executor_options[ "max_size"] configuration.set(executor_section, "max_size", max_size) executor_metadata = { "executor_name": executor_name, "args": {param: False for param in false_params}, } executor_metadata["args"]["out"] = True test_config.executors.append(executor_metadata) tmp_default_config.save() # Init and register it dispatcher = Dispatcher(test_config.client.session, tmp_default_config.config_file_path) selected_workspace = random.choice(workspaces) print(selected_workspace) ws_responses = deepcopy(executor_options["ws_responses"]) run_data = deepcopy(executor_options["data"]) if "workspace" in run_data: run_data["workspace"] = run_data["workspace"].format( selected_workspace) test_config.ws_data = {"run_data": run_data, "ws_responses": ws_responses} await dispatcher.register(test_config.registration_token) await dispatcher.connect() history = test_logger_handler.history assert len(test_config.ws_data["ws_responses"]) == 0 for log in executor_options["logs"]: min_count = 1 if "min_count" not in log else log["min_count"] max_count = sys.maxsize if "max_count" not in log else log["max_count"] assert (max_count >= len( list( filter( lambda x: x.levelname == log["levelname"] and log["msg"] in x.message, history, ))) >= min_count), log["msg"]
def test_execute_agent(): session = Session() # Initial set up res = session.post(api_url(HOST, API_PORT, postfix='/_api/login'), json={ 'email': USER, 'password': PASS }) assert res.status_code == 200, res.text session_res = session.get(api_url(HOST, API_PORT, postfix='/_api/session')) res = session.post(api_url(HOST, API_PORT, postfix='/_api/v2/ws/'), json={'name': WORKSPACE}) assert res.status_code == 201, res.text res = session.get(api_url(HOST, API_PORT, postfix='/_api/v2/agent_token/')) assert res.status_code == 200, res.text token = res.json()['token'] # Config set up config.set(Sections.TOKENS, "registration", token) config.remove_option(Sections.TOKENS, "agent") config.set(Sections.SERVER, "workspace", WORKSPACE) config.set(Sections.AGENT, "agent_name", AGENT_NAME) config.set(Sections.AGENT, "executors", EXECUTOR_NAME) path_to_basic_executor = (Path(__file__).parent.parent.parent / 'data' / 'basic_executor.py') executor_section = Sections.EXECUTOR_DATA.format(EXECUTOR_NAME) params_section = Sections.EXECUTOR_PARAMS.format(EXECUTOR_NAME) for section in [executor_section, params_section]: if section not in config: config.add_section(section) config.set(Sections.EXECUTOR_DATA.format(EXECUTOR_NAME), "cmd", f"python {path_to_basic_executor}") config.set(params_section, "out", "True") [ config.set(params_section, param, "False") for param in ["count", "spare", "spaced_before", "spaced_middle", "err", "fails"] ] save_config(CONFIG_DIR) # Init dispatcher! command = [ 'faraday-dispatcher', f'--config-file={CONFIG_DIR}', f'--logdir={LOGGER_DIR}' ] p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) time.sleep(2) # If fails check time # Checking dispatcher connection res = session.get( api_url(HOST, API_PORT, postfix=f'/_api/v2/ws/{WORKSPACE}/agents/')) assert res.status_code == 200, res.text res_data = res.json() assert len(res_data) == 1, p.communicate(timeout=0.1) agent = res_data[0] agent_id = agent["id"] if agent_ok_status_keys_set != set(agent.keys()): print( "Keys set from agent endpoint differ from expected ones, checking if its a superset" ) assert agent_ok_status_keys_set.issubset(set(agent.keys())) for key in agent_ok_status_dict: assert agent[key] == agent_ok_status_dict[key], [ agent, agent_ok_status_dict ] # Run executor! res = session.post(api_url( HOST, API_PORT, postfix=f'/_api/v2/ws/{WORKSPACE}/agents/{agent["id"]}/run/'), json={ 'csrf_token': session_res.json()['csrf_token'], 'executorData': { "agent_id": agent_id, "executor": EXECUTOR_NAME, "args": { "out": "json" } } }) assert res.status_code == 200, res.text time.sleep(2) # If fails check time # Test results res = session.get( api_url(HOST, API_PORT, postfix=f'/_api/v2/ws/{WORKSPACE}/hosts')) host_dict = res.json() assert host_dict["total_rows"] == 1, (res.text, host_dict) host = host_dict["rows"][0]["value"] for key in host_data: if key == "hostnames": assert set(host[key]) == set(host_data[key]) else: assert host[key] == host_data[key] assert host["vulns"] == 1 res = session.get( api_url(HOST, API_PORT, postfix=f'/_api/v2/ws/{WORKSPACE}/vulns')) vuln_dict = res.json() assert vuln_dict["count"] == 1 vuln = vuln_dict["vulnerabilities"][0]["value"] for key in vuln_data: if key == 'impact': for k_key in vuln['impact']: if k_key in vuln_data['impact']: assert vuln['impact'][k_key] == vuln_data['impact'][k_key] else: assert not vuln['impact'][k_key] else: assert vuln[key] == vuln_data[key] assert vuln["target"] == host_data['ip']
def test_execute_agent(): session = Session() # Initial set up res = session.post( api_url(HOST, API_PORT, postfix="/_api/login"), json={ "email": USER, "password": PASS }, ) assert res.status_code == 200, res.text # session_res = session.get(api_url(HOST, API_PORT, postfix="/_api/session")) res = session.post(api_url(HOST, API_PORT, postfix="/_api/v2/ws/"), json={"name": WORKSPACE}) assert res.status_code == 201, res.text res = session.get(api_url(HOST, API_PORT, postfix="/_api/v2/agent_token/")) assert res.status_code == 200, res.text token = res.json()["token"] # Config set up config.set(Sections.TOKENS, "registration", token) config.remove_option(Sections.TOKENS, "agent") config.set(Sections.SERVER, "workspaces", WORKSPACE) config.set(Sections.SERVER, "ssl", SSL) config.set(Sections.AGENT, "agent_name", AGENT_NAME) config.set(Sections.AGENT, "executors", EXECUTOR_NAME) path_to_basic_executor = Path( __file__).parent.parent.parent / "data" / "basic_executor.py" executor_section = Sections.EXECUTOR_DATA.format(EXECUTOR_NAME) params_section = Sections.EXECUTOR_PARAMS.format(EXECUTOR_NAME) for section in [executor_section, params_section]: if section not in config: config.add_section(section) config.set( Sections.EXECUTOR_DATA.format(EXECUTOR_NAME), "cmd", f"python {path_to_basic_executor}", ) config.set(params_section, "out", "True") [ config.set(params_section, param, "False") for param in [ "count", "spare", "spaced_before", "spaced_middle", "err", "fails", ] ] with tempfile.TemporaryDirectory() as tempdirfile: config_pathfile = Path(tempdirfile) save_config(config_pathfile) # Init dispatcher! command = [ "faraday-dispatcher", "run", f"--config-file={config_pathfile}", f"--logdir={LOGGER_DIR}", ] p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) time.sleep(2) # If fails check time # Checking dispatcher connection res = session.get( api_url(HOST, API_PORT, postfix=f"/_api/v2/ws/{WORKSPACE}/agents/")) assert res.status_code == 200, res.text res_data = res.json() assert len(res_data) == 1, p.communicate(timeout=0.1) agent = res_data[0] agent_id = agent["id"] if agent_ok_status_keys_set != set(agent.keys()): print( "Keys set from agent endpoint differ from expected ones, checking if its a superset" ) print(f"agent_ok_status_keys_set= {agent_ok_status_keys_set}") print(f"agent.keys() = {agent.keys()}") assert agent_ok_status_keys_set.issubset(set(agent.keys())) for key in agent_ok_status_dict: assert agent[key] == agent_ok_status_dict[key], [ agent, agent_ok_status_dict, ] # Run executor! res = session.post( api_url( HOST, API_PORT, postfix=f'/_api/v2/ws/{WORKSPACE}/agents/{agent["id"]}/run/', ), json={ # "csrf_token": session_res.json()["csrf_token"], "executorData": { "agent_id": agent_id, "executor": EXECUTOR_NAME, "args": { "out": "json" }, }, }, ) assert res.status_code == 200, res.text command_id = res.json()["command_id"] # Command ID should be in progress! res = session.get( api_url( HOST, API_PORT, postfix=f"/_api/v2/ws/{WORKSPACE}/commands/{command_id}/", ), ) assert res.status_code == 200, res.text command_check_response = res.json() assert command_check_response["import_source"] == "agent" assert command_check_response["creator"] is None assert command_check_response["duration"] == "In progress" time.sleep(2) # If fails check time # Command ID should not be in progress! res = session.get( api_url( HOST, API_PORT, postfix=f"/_api/v2/ws/{WORKSPACE}/commands/{command_id}/", ), ) assert res.status_code == 200, res.text command_check_response = res.json() assert command_check_response["duration"] != "In progress" # Test results res = session.get( api_url(HOST, API_PORT, postfix=f"/_api/v2/ws/{WORKSPACE}/hosts")) host_dict = res.json() assert host_dict["count"] == 1, (res.text, host_dict) host = host_dict["rows"][0]["value"] for key in host_data: if key == "hostnames": assert set(host[key]) == set(host_data[key]) else: assert host[key] == host_data[key] assert host["vulns"] == 1 res = session.get( api_url(HOST, API_PORT, postfix=f"/_api/v2/ws/{WORKSPACE}/vulns")) vuln_dict = res.json() assert vuln_dict["count"] == 1 vuln = vuln_dict["vulnerabilities"][0]["value"] for key in vuln_data: if key == "impact": for k_key in vuln["impact"]: if k_key in vuln_data["impact"]: assert vuln["impact"][k_key] == vuln_data["impact"][ k_key] else: assert not vuln["impact"][k_key] else: assert vuln[key] == vuln_data[key] assert vuln["target"] == host_data["ip"]
async def register(self, registration_token=None): if not await self.check_connection(): exit(1) if self.agent_token is None: try: control_registration_token("token", registration_token) except ValueError as ex: print(f"{Bcolors.FAIL}{ex.args[0]}") logger.error(ex.args[0]) exit(1) token_registration_url = api_url( self.host, self.api_port, postfix="/_api/v3/agent_registration", secure=self.api_ssl_enabled, ) logger.info(f"token_registration_url: {token_registration_url}") try: token_response = await self.session.post( token_registration_url, json={ "token": registration_token, "name": self.agent_name, "workspaces": self.workspaces, }, **self.api_kwargs, ) token = await token_response.json() self.agent_token = token["token"] if Sections.TOKENS not in config: config.add_section(Sections.TOKENS) config.set(Sections.TOKENS, "agent", self.agent_token) save_config(self.config_path) except ClientResponseError as e: if e.status == 404: logger.error( "404 HTTP ERROR received: Can't connect to the server") elif e.status == 401: logger.error( "Invalid registration token, please reset and retry. " "If the error persist, you should try to edit the " "registration token with the wizard command " "`faraday-dispatcher config-wizard`\nHint: If the faraday " "version is not the expected this could fail, check " "https://github.com/infobyte/faraday_agent_dispatcher/blob/master/RELEASE.md" ) else: logger.info(f"Unexpected error: {e}") logger.debug(msg="Exception raised", exc_info=e) exit(1) except ClientConnectorError as e: logger.debug(msg="Connection con error failed", exc_info=e) logger.error("Can connect to server") try: self.websocket_token = await self.reset_websocket_token() logger.info("Registered successfully") except ClientResponseError as e: error_msg = ("Invalid agent token, please reset and retry. If " "the error persist, you should remove the agent " "token with the wizard command `faraday-dispatcher " "config-wizard`") logger.error(error_msg) self.agent_token = None logger.debug(msg="Exception raised", exc_info=e) exit(1)