def setUp(self):
super(TestChat, self).setUp()
self.chat = ChatFactory()
self.agent = AgentFactory()
self.question = "text question"
self.answer = "text_answer"
self.message_id = self._gen_id()
def setUp(self):
super(TestColdCenter, self).setUp()
self.agent = AgentFactory()
self.agent.save()
self.asker = AskerFactory()
self.asker.save()
self.question = "question content"
self.answer = "answer content"
self.message_id = self._gen_id()
def test_invalid_body(self, test_client, session):
agent = AgentFactory.create(workspace=self.workspace)
session.add(agent)
session.commit()
res = test_client.post(self.url() + f'{agent.id}/run/',
data='[" broken]"{')
assert res.status_code == 400
def setUp(self):
super(TestChat, self).setUp()
self.chat = ChatFactory()
self.agent = AgentFactory()
self.question = "text question"
self.answer = "text_answer"
self.message_id = self._gen_id()
def test_run_fails(self, test_client, session,csrf_token):
workspace = WorkspaceFactory.create()
session.add(workspace)
other_workspace = WorkspaceFactory.create()
session.add(other_workspace)
session.commit()
agent = AgentFactory.create(
workspaces=[workspace, other_workspace]
)
executor = ExecutorFactory.create(agent=agent)
session.add(executor)
session.commit()
payload = {
'csrf_token': csrf_token,
'executorData': {
"args": {
"param1": True
},
"executor": executor.name
},
}
res = test_client.post(
self.url(agent) + 'run/',
json=payload
)
assert res.status_code == 404
def test_happy_path_valid_json(self, test_client, session, csrf_token):
agent = AgentFactory.create(workspaces=[self.workspace])
executor = ExecutorFactory.create(agent=agent)
executor2 = ExecutorFactory.create(agent=agent)
session.add(executor)
session.commit()
assert agent.last_run is None
assert executor.last_run is None
assert executor2.last_run is None
payload = {
'csrf_token': csrf_token,
'executorData': {
"args": {
"param_name": "test"
},
"executor": executor.name
},
}
res = test_client.post(urljoin(self.url(agent), 'run'), json=payload)
assert res.status_code == 200
command_id = res.json["command_id"]
command = Command.query.filter(
Command.workspace_id == self.workspace.id).one()
assert command_id == command.id
assert agent.last_run is not None
assert executor.last_run is not None
assert executor2.last_run is None
assert agent.last_run == executor.last_run
def test_get_workspaced_other_fails(self, test_client, session):
other_workspace = WorkspaceFactory.create()
session.add(other_workspace)
agent = AgentFactory.create(workspaces=[other_workspace], active=True)
session.commit()
res = test_client.get(self.url(agent))
assert res.status_code == 404
def test_invalid_body(self, test_client, session):
agent = AgentFactory.create(workspaces=[self.workspace])
session.add(agent)
session.commit()
res = test_client.post(urljoin(self.url(agent), 'run'),
data='[" broken]"{')
assert res.status_code == 400
def test_update_bug_case(self, test_client, session):
agent = AgentFactory.create(workspace=self.workspace)
session.add(agent)
session.commit()
update_data = {"id": 1, "name": "Agent test", "is_online": True}
res = test_client.put(self.url(agent.id), data=update_data)
assert res.status_code == 200
def test_update_agent(self, test_client, session):
agent = AgentFactory.create(workspace=self.workspace, active=True)
session.commit()
raw_agent = self.create_raw_agent(active=False)
res = test_client.put(self.url(agent.id), data=raw_agent)
assert res.status_code == 200
assert not res.json['active']
def test_get_workspaced(self, test_client, session):
workspace = WorkspaceFactory.create()
session.add(workspace)
agent = AgentFactory.create(workspaces=[self.workspace], active=True)
session.commit()
res = test_client.get(self.url(agent))
assert res.status_code == 200
assert 'workspaces' not in res.json
def test_delete_agent(self, test_client, session):
initial_agent_count = len(session.query(Agent).all())
agent = AgentFactory.create(workspace=self.workspace)
session.commit()
assert len(session.query(Agent).all()) == initial_agent_count + 1
res = test_client.delete(self.url(agent.id))
assert res.status_code == 204
assert len(session.query(Agent).all()) == initial_agent_count
def test_run_agent_invalid_missing_executorData(self, csrf_token, session,
test_client):
agent = AgentFactory.create(workspaces=[self.workspace])
session.add(agent)
session.commit()
payload = {'csrf_token': csrf_token}
res = test_client.post(self.url(agent) + 'run/', json=payload)
assert res.status_code == 400
def setUp(self):
super(TestColdCenter, self).setUp()
self.agent = AgentFactory()
self.agent.save()
self.asker = AskerFactory()
self.asker.save()
self.question = "question content"
self.answer = "answer content"
self.message_id = self._gen_id()
def _join_agent(test_client, session):
agent = AgentFactory.create(token='pepito')
session.add(agent)
session.commit()
headers = {"Authorization": "Agent {}".format(agent.token)}
token = test_client.post('v2/agent_websocket_token/',
headers=headers).json['token']
return token
def _join_agent(self, test_client, session):
agent = AgentFactory.create(token='pepito')
session.add(agent)
session.commit()
headers = {"Authorization": f"Agent {agent.token}"}
token = test_client.post(self.check_url('/v2/agent_websocket_token/'),
headers=headers).json['token']
return token
def test_get_not_workspaced(self, test_client, session):
workspace = WorkspaceFactory.create()
session.add(workspace)
agent = AgentFactory.create(workspaces=[workspace], active=True)
session.commit()
res = test_client.get(self.url(agent))
assert res.status_code == 200
assert len(res.json['workspaces']) == 1
assert workspace.name in res.json['workspaces'][0]
def test_new_executors_not_in_database(self, session):
agent = AgentFactory.create()
executors = [{
'executor_name': 'nmap_executor',
'args': {
'param1': True
}
}]
assert update_executors(agent, executors)
def test_run_agent(self, session, csrf_token, test_client):
agent = AgentFactory.create(workspaces=[self.workspace])
session.add(agent)
session.commit()
payload = {
'csrf_token': csrf_token,
'executorData': '',
}
res = test_client.post(urljoin(self.url(agent), 'run'), json=payload)
assert res.status_code == 400
def test_run_agent(self, session, csrf_token, test_client):
agent = AgentFactory.create(workspace=self.workspace)
session.add(agent)
session.commit()
payload = {
'csrf_token': csrf_token,
'executorData': '',
}
res = test_client.post(self.url() + f'{agent.id}/run/', json=payload)
assert res.status_code == 400
def test_invalid_json_on_executorData_breaks_the_api(
self, csrf_token, session, test_client):
agent = AgentFactory.create(workspaces=[self.workspace])
session.add(agent)
session.commit()
payload = {
'csrf_token': csrf_token,
'executorData': '[][dassa',
}
res = test_client.post(self.url(agent) + 'run/', json=payload)
assert res.status_code == 400
def test_update_agent(self, test_client, session):
workspace = WorkspaceFactory.create()
session.add(workspace)
agent = AgentFactory.create(workspaces=[workspace], active=True)
session.commit()
raw_agent = self.create_raw_agent(active=False, workspaces=[workspace])
res = test_client.put(self.url(agent.id), data=raw_agent)
assert res.status_code == 200, (res.json, raw_agent)
assert not res.json['active']
assert len(res.json['workspaces']) == 1
assert workspace.name in res.json['workspaces'][0]
def test_remove_all_executors(self, session):
agent = AgentFactory.create()
params = {'old_param': True}
executor = ExecutorFactory.create(name='old_executor',
agent=agent,
parameters_metadata=params)
session.add(executor)
session.commit()
executors = []
assert update_executors(agent, executors)
count_executors = Executor.query.filter_by(agent=agent).count()
assert count_executors == 0
def test_executors_with_some_invalid_executors(self, session):
agent = AgentFactory.create()
executors = [{
'invalid_key': 'nmap_executor'
}, {
'executor_name': 'executor 1',
'args': {
'param1': True
}
}]
assert update_executors(agent, executors)
count_executors = Executor.query.filter_by(agent=agent).count()
assert count_executors == 1
def test_invalid_executor(self, test_client, session, csrf_token):
agent = AgentFactory.create(workspaces=[self.workspace])
session.add(agent)
session.commit()
payload = {
'csrf_token': csrf_token,
'executorData': {
"args": {
"param1": True
},
"executor": "executor_name"
},
}
res = test_client.post(self.url(agent) + 'run/', json=payload)
assert res.status_code == 400
def test_happy_path_valid_json(self, test_client, session, csrf_token):
agent = AgentFactory.create(workspace=self.workspace)
session.add(agent)
session.commit()
payload = {
'csrf_token': csrf_token,
'executorData': {
"args": {
"param1": True
},
"executor": "executor_name"
},
}
res = test_client.post(self.url() + f'{agent.id}/run/', json=payload)
assert res.status_code == 200
def test_update_bug_case(self, test_client, session):
workspace = WorkspaceFactory.create()
session.add(workspace)
agent = AgentFactory.create(workspaces=[workspace])
session.add(agent)
session.commit()
update_data = {
"id": 1,
"name": "Agent test",
"is_online": True,
"workspaces": [workspace.name]
}
res = test_client.put(self.url(agent.id), data=update_data)
assert res.status_code == 200, (res.json, update_data)
assert workspace.name in res.json['workspaces']
assert len(res.json['workspaces']) == 1
def test_update_agent_add_a_inexistent_workspace(self, test_client,
session):
workspace = WorkspaceFactory.create()
session.add(workspace)
other_workspace = WorkspaceFactory.create()
session.add(other_workspace)
agent = AgentFactory.create(workspaces=[workspace], active=True)
session.commit()
raw_agent = self.create_raw_agent(
workspaces=[workspace, other_workspace])
raw_agent["workspaces"] = ["donotexist"]
res = test_client.put(self.url(agent.id), data=raw_agent)
assert res.status_code == 404
workspaces = Agent.query.get(agent.id).workspaces
assert len(workspaces) == 1
assert workspace in workspaces
def test_executor_already_in_database_but_new_parameters_incoming(
self, session):
agent = AgentFactory.create()
old_params = {'old_param': True}
executor = ExecutorFactory.create(agent=agent,
parameters_metadata=old_params)
session.add(executor)
session.commit()
new_params = old_params
new_params.update({'param1': True})
executors = [{'executor_name': executor.name, 'args': new_params}]
assert update_executors(agent, executors)
from_db_executor = Executor.query.filter_by(id=executor.id,
agent=agent).first()
assert from_db_executor.parameters_metadata == new_params
def test_happy_path_valid_json(self, test_client, session, csrf_token):
agent = AgentFactory.create(workspaces=[self.workspace])
executor = ExecutorFactory.create(agent=agent)
session.add(executor)
session.commit()
payload = {
'csrf_token': csrf_token,
'executorData': {
"args": {
"param1": True
},
"executor": executor.name
},
}
res = test_client.post(self.url(agent.id) + 'run/', json=payload)
assert res.status_code == 200
def test_update_agent_delete_a_workspace(self, test_client, session):
workspace = WorkspaceFactory.create()
session.add(workspace)
other_workspace = WorkspaceFactory.create()
session.add(other_workspace)
agent = AgentFactory.create(workspaces=[workspace, other_workspace],
active=True)
session.commit()
raw_agent = self.create_raw_agent(workspaces=[workspace])
res = test_client.put(self.url(agent.id), data=raw_agent)
assert res.status_code == 200
assert len(res.json['workspaces']) == 1
assert other_workspace.name not in res.json['workspaces']
assert workspace.name in res.json['workspaces']
workspaces = Agent.query.get(agent.id).workspaces
assert len(workspaces) == 1
assert workspaces[0] == workspace
def test_workspaced_delete(self, session, test_client):
initial_agent_count = len(session.query(Agent).all())
other_workspace = WorkspaceFactory.create()
session.add(other_workspace)
agent = AgentFactory.create(
workspaces=[self.workspace, other_workspace])
session.commit()
assert len(session.query(Agent).all()) == initial_agent_count + 1
res = test_client.delete(self.url(agent.id))
assert res.status_code == 204
assert len(session.query(Agent).all()) == initial_agent_count + 1
res = test_client.delete(self.url(agent.id))
assert res.status_code == 404
res = test_client.get(self.url(agent.id))
assert res.status_code == 404
workspaces = Agent.query.get(agent.id).workspaces
assert len(workspaces) == 1
assert other_workspace in workspaces
class TestChat(BaseTestColdCenter):
def setUp(self):
super(TestChat, self).setUp()
self.chat = ChatFactory()
self.agent = AgentFactory()
self.question = "text question"
self.answer = "text_answer"
self.message_id = self._gen_id()
def test_creation_state(self):
self.assertChatState(self.chat, Chat.PENDING)
def test_handle_message_no_agent_to_assign(self):
self.assertRaises(NoAgentFound, self.chat.handle_message, self.message_id, self.chat.asker.id_service,
self.question, self.listenclosely_app)
def test_handle_message_assign_agent(self):
self._register(self.agent)
self.agent.save()
self.chat.handle_message(self.message_id, self.chat.asker.id_service, self.question, self.listenclosely_app)
self.chat.save()
self.assertChatState(self.chat, Chat.LIVE)
self.assertEqual(self.chat.agent, self.agent)
self.assertAgentState(self.chat.agent, Agent.BUSY)
self.assertEqual(1, Message.objects.count())
message = Message.objects.all()[0]
self.assertEqual(message.id_service_in, self.message_id)
self.assertEqual(message.content, self.question)
self.assertEqual(message.type, Message.INCOMING)
self.assertEqual(message.chat, self.chat)
self.assertMessageServiceBackend(message.id_service_out, self.chat.agent.id_service, self.question)
def test_handle_message_already_assigned(self):
self._register(self.agent)
self.agent.save()
self._attend(self.agent, self.chat)
self.agent.save()
self.assertEqual(self.chat.agent, self.agent)
self.assertAgentState(self.agent, Agent.BUSY)
self.chat.handle_message(self.message_id, self.chat.asker.id_service, self.question, self.listenclosely_app)
self.chat.save()
self.assertChatState(self.chat, Chat.LIVE)
self.assertEqual(1, Message.objects.count())
message = Message.objects.all()[0]
self.assertEqual(message.id_service_in, self.message_id)
self.assertEqual(message.content, self.question)
self.assertEqual(message.type, Message.INCOMING)
self.assertEqual(message.chat, self.chat)
self.assertMessageServiceBackend(message.id_service_out, self.chat.agent.id_service, self.question)
def test_handle_message_answer(self):
self._register(self.agent)
self.agent.save()
self._attend(self.agent, self.chat)
self.agent.save()
self.assertEqual(self.chat.agent, self.agent)
self.assertAgentState(self.agent, Agent.BUSY)
self.chat.handle_message(self.message_id, self.chat.agent.id_service, self.answer, self.listenclosely_app)
self.chat.save()
self.assertChatState(self.chat, Chat.LIVE)
self.assertEqual(1, Message.objects.count())
message = Message.objects.all()[0]
self.assertEqual(message.id_service_in, self.message_id)
self.assertEqual(message.content, self.answer)
self.assertEqual(message.type, Message.OUTGOING)
self.assertEqual(message.chat, self.chat)
self.assertMessageServiceBackend(message.id_service_out, self.chat.asker.id_service, self.answer)
def test_terminate(self):
self._register(self.agent)
self.agent.save()
self.chat.handle_message(self.message_id, self.chat.asker.id_service, self.question, self.listenclosely_app)
self.chat.save()
self.assertChatState(self.chat, Chat.LIVE)
self.assertAgentState(self.chat.agent, Agent.BUSY)
self.chat.terminate()
self.assertChatState(self.chat, Chat.TERMINATED)
self.assertAgentState(self.chat.agent, Agent.ONLINE)
class TestColdCenter(BaseTestColdCenter):
def setUp(self):
super(TestColdCenter, self).setUp()
self.agent = AgentFactory()
self.agent.save()
self.asker = AskerFactory()
self.asker.save()
self.question = "question content"
self.answer = "answer content"
self.message_id = self._gen_id()
def test_attend_pending_one_query_one_free_agent(self):
self.query = ChatFactory()
self.query.asker = self.asker
self.query.save()
self._register(self.agent)
self.agent.save()
self.assertNumberAgentEachState(0, 1, 0)
self.assertNumberChatEachState(1, 0, 0)
queries_attended = self.listenclosely_app.attend_pendings()
self.assertEqual(1, len(queries_attended))
self.assertEqual(queries_attended[0], Chat.live.all()[0])
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
self.query = Chat.objects.all()[0]
self.assertChatState(self.query, Chat.LIVE)
self.assertEqual(self.agent, self.query.agent)
self.assertAgentState(self.query.agent, Agent.BUSY)
def test_attend_pending_one_query_no_free_agent(self):
self.query = ChatFactory()
self.query.asker = self.asker
self.query.save()
self.assertNumberAgentEachState(1, 0, 0)
self.assertNumberChatEachState(1, 0, 0)
queries_attended = self.listenclosely_app.attend_pendings()
self.assertEqual(0, len(queries_attended))
self.assertNumberAgentEachState(1, 0, 0)
self.assertNumberChatEachState(1, 0, 0)
def test_attend_pending_no_peding_queries(self):
self.query = ChatFactory()
self._register(self.agent)
self.agent.save()
self.query.handle_message(self.message_id, self.query.asker.id_service, self.question, self.listenclosely_app)
self.query.save()
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
queries_attended = self.listenclosely_app.attend_pendings()
self.assertEqual(0, len(queries_attended))
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
def test_terminate_obsolete_one(self):
self.query = ChatFactory()
self._register(self.agent)
self.agent.save()
self.query.handle_message(self.message_id, self.query.asker.id_service, self.question, self.listenclosely_app)
self.query.save()
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
self.listenclosely_app.time_obsolete_offset = -1
queries_terminated = self.listenclosely_app.terminate_obsolete()
self.assertEqual(1, len(queries_terminated))
self.assertEqual(queries_terminated[0], Chat.terminated.all()[0])
self.assertNumberChatEachState(0, 0, 1)
self.assertNumberAgentEachState(0, 1, 0)
def test_terminate_obsolete_no_obsolete(self):
self.query = ChatFactory()
self._register(self.agent)
self.agent.save()
self.query.handle_message(self.message_id, self.query.asker.id_service, self.question, self.listenclosely_app)
self.query.save()
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
queries_terminated = self.listenclosely_app.terminate_obsolete()
self.assertEqual(0, len(queries_terminated))
self.assertNumberChatEachState(0, 1, 0)
self.assertNumberAgentEachState(0, 0, 1)
def test_terminate_obsolete_no_live(self):
self.query = ChatFactory()
self.assertNumberChatEachState(1, 0, 0)
self.listenclosely_app.time_obsolete_offset = -1
queries_terminated = self.listenclosely_app.terminate_obsolete()
self.assertEqual(0, len(queries_terminated))
self.assertNumberChatEachState(1, 0, 0)
def test_on_message_asker_new_question_with_free_agent(self):
self._register(self.agent)
self.agent.save()
self.assertNumberAgentEachState(0, 1, 0)
self.assertNumberChatEachState(0, 0, 0)
self.listenclosely_app.on_message(self.message_id, self.asker.id_service, self.question)
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
message = Chat.live.all()[0].messages.all()[0]
self.assertMessageServiceBackend(message.id_service_out, self.agent.id_service, self.question)
def test_on_message_new_asker_new_question_with_free_agent(self):
self._register(self.agent)
self.agent.save()
self.assertNumberAgentEachState(0, 1, 0)
self.assertNumberChatEachState(0, 0, 0)
new_asker_id = "new_asker_id"
self.listenclosely_app.on_message(self.message_id, new_asker_id, self.question)
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
asker = Asker.objects.get(id_service=new_asker_id)
self.assertEqual(asker.id_service, new_asker_id)
message = Chat.live.all()[0].messages.all()[0]
self.assertMessageServiceBackend(message.id_service_out, self.agent.id_service, self.question)
def test_on_message_with_no_free_agent(self):
self.assertNumberAgentEachState(1, 0, 0)
self.assertNumberChatEachState(0, 0, 0)
self.listenclosely_app.on_message(self.message_id, self.asker.id_service, self.question)
self.assertNumberAgentEachState(1, 0, 0)
self.assertNumberChatEachState(1, 0, 0)
self.query = Chat.pending.all()[0]
self.assertEqual(1, self.query.messages.count())
self.assertEqual(None, self.query.messages.all()[0].t_sent)
def test_on_message_from_agent_with_no_query(self):
self._register(self.agent)
self.agent.save()
self.assertNumberAgentEachState(0, 1, 0)
self.assertNumberChatEachState(0, 0, 0)
self.listenclosely_app.on_message(self.message_id, self.agent.id_service, self.answer)
self.assertNumberAgentEachState(0, 1, 0)
self.assertNumberChatEachState(0, 0, 0)
self.assertEqual(0, Message.objects.count())
def test_on_message_answer_to_current_query(self):
self._register(self.agent)
self.agent.save()
self.assertNumberAgentEachState(0, 1, 0)
self.assertNumberChatEachState(0, 0, 0)
new_asker_id = "new_asker_id"
self.listenclosely_app.on_message(self.message_id, new_asker_id, self.question)
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
message = Chat.live.all()[0].messages.all()[0]
self.assertMessageServiceBackend(message.id_service_out, self.agent.id_service, self.question)
self.listenclosely_app.on_message(self.message_id + "2", self.agent.id_service, self.answer)
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
query = Chat.live.get(agent__id_service=self.agent.id_service)
self.assertEqual(2, query.messages.count())
message = Chat.live.all()[0].messages.all()[1]
self.assertMessageServiceBackend(message.id_service_out, new_asker_id, self.answer)
def test_on_message_question_to_current_attended_query(self):
self._register(self.agent)
self.agent.save()
self.assertNumberAgentEachState(0, 1, 0)
self.assertNumberChatEachState(0, 0, 0)
self.listenclosely_app.on_message(self.message_id, self.asker.id_service, self.question)
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
message = Chat.live.all()[0].messages.all()[0]
self.assertMessageServiceBackend(message.id_service_out, self.agent.id_service, self.question)
self.listenclosely_app.on_message(self.message_id + "2", self.asker.id_service, self.question + "2")
self.assertNumberAgentEachState(0, 0, 1)
self.assertNumberChatEachState(0, 1, 0)
message = Chat.live.all()[0].messages.all()[1]
self.assertMessageServiceBackend(message.id_service_out, self.agent.id_service, self.question + "2")