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 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_start_with_bad_config(test_config: FaradayTestConfig, tmp_default_config):
# Config
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", "NotOk" * 5)
configuration.set(Sections.EXECUTOR_DATA.format("ex1"), "cmd", 'exit 1')
tmp_default_config.save()
# Init and register it
dispatcher = Dispatcher(test_config.client.session, tmp_default_config.config_file_path)
with pytest.raises(AssertionError):
await dispatcher.register()
async def test_start_and_register(test_config: FaradayTestConfig, tmp_default_config):
# Config
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.EXECUTOR_DATA.format("ex1"), "cmd", 'exit 1')
tmp_default_config.save()
# Init and register it
dispatcher = Dispatcher(test_config.client.session, tmp_default_config.config_file_path)
await dispatcher.register()
# 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
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 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_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_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_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
if Sections.TOKENS not in configuration:
configuration.add_section(Sections.TOKENS)
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(test_config.registration_token)
await dispatcher.connect()
assert len(test_config.ws_data["ws_responses"]) == 0
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)