async def register(self): if self.agent_token is None: registration_token = self.agent_token = config.get( Sections.TOKENS, "registration") assert registration_token is not None, "The registration token is mandatory" token_registration_url = api_url( self.host, self.api_port, postfix=f"/_api/v2/ws/{self.workspace}/agent_registration/") 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 }) assert token_response.status == 201 token = await token_response.json() self.agent_token = token["token"] config.set(Sections.TOKENS, "agent", self.agent_token) save_config(self.config_path) except ClientResponseError as e: if e.status == 404: logger.info( f'404 HTTP ERROR received: Workspace "{self.workspace}" not found' ) return else: logger.info(f"Unexpected error: {e}") raise e self.websocket_token = await self.reset_websocket_token()
def post_url(): host = config.get('server', 'host') port = config.get('server', 'api_port') return api_url( host, port, postfix= f"/_api/v2/ws/{config.get('server', 'workspace')}/bulk_create/")
def post_url(self): host = config.get("server", "host") port = config.get("server", "api_port") return api_url( host, port, postfix="/_api/v3/ws/" f"{self.workspace}/bulk_create", secure=self.api_ssl_enabled, )
async def reset_websocket_token(self): # I'm built so I ask for websocket token headers = {"Authorization": f"Agent {self.agent_token}"} websocket_token_response = await self.session.post( api_url(self.host, self.api_port, postfix='/_api/v2/agent_websocket_token/', secure=self.api_ssl_enabled), headers=headers, **self.api_kwargs ) websocket_token_json = await websocket_token_response.json() return websocket_token_json["token"]
async def register(self): if self.agent_token is None: registration_token = self.agent_token = config.get(Sections.TOKENS, "registration") assert registration_token is not None, "The registration token is mandatory" token_registration_url = api_url(self.host, self.api_port, postfix=f"/_api/v2/ws/{self.workspace}/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}, **self.api_kwargs ) token = await token_response.json() self.agent_token = token["token"] config.set(Sections.TOKENS, "agent", self.agent_token) save_config(self.config_path) except ClientResponseError as e: if e.status == 404: logger.error(f'404 HTTP ERROR received: Workspace "{self.workspace}" not found') 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`") else: logger.info(f"Unexpected error: {e}") logger.debug(msg="Exception raised", exc_info=e) return 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 " \ f"the agent token with the wizard command `faraday-dispatcher " \ f"config-wizard`" logger.error(error_msg) self.agent_token = None logger.debug(msg="Exception raised", exc_info=e) return
async def check_connection(self): server_url = api_url(self.host, self.api_port, postfix="/_api/v3/info", secure=self.api_ssl_enabled) logger.debug(f"Validating server connection with {server_url}") try: kwargs = self.api_kwargs.copy() if "DISPATCHER_TEST" in os.environ and os.environ[ "DISPATCHER_TEST"] == "True": kwargs["timeout"] = ClientTimeout(total=1) # > The below code allows this get to be canceled, # > But breaks only ours Gitlab CI tests (local OK) # check_connection_task = asyncio.create_task( # self.session.get(server_url, **kwargs) # ) # self.executor_tasks[Dispatcher.TaskLabels.CONNECTION_CHECK].\ # append(check_connection_task) # await check_connection_task await self.session.get(server_url, **kwargs) except (ClientConnectorCertificateError, ClientConnectorSSLError) as e: logger.debug("Invalid SSL Certificate", exc_info=e) print(f"{Bcolors.FAIL}Invalid SSL Certificate, use " "`faraday-dispatcher config-wizard` and check the " "certificate configuration") return False except ClientConnectorError as e: logger.error("Can not connect to Faraday server") logger.debug("Connect failed traceback", exc_info=e) return False except asyncio.TimeoutError as e: logger.error("Faraday server timed-out. " "TIP: Check ssl configuration") logger.debug("Timeout error. Check ssl", exc_info=e) return False except asyncio.CancelledError: logger.info("User sent close signal") return False return True
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.instance: config.instance[Sections.TOKENS] = {} config.instance[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: if e.status == 402: error_msg = "Unauthorized. Is your license expired or invalid?" else: 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)
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"]
def post_url(self): host = config.get('server', 'host') port = config.get('server', 'api_port') return api_url(host, port, postfix=f"/_api/v2/ws/{config.get('server', 'workspace')}/bulk_create/", secure=self.api_ssl_enabled)