def test_get_task_status_no_filter(mock_task_db):
"""Check system's tasks (No filters)
"""
result = task.get_task_status()
cur = get_fake_database_data('schema_tasks_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT task_id) FROM tasks")
rows = cur.fetchone()
assert result.total_affected_items == rows[0]
def test_check_total_items(mock_wdb):
"""Test the number of returned items."""
# load test database and make the query
cur = get_fake_database_data('schema_mitre_test.sql').cursor()
cur.execute(f'SELECT COUNT(DISTINCT id) FROM attack')
rows = cur.fetchone()
expected_total_items = rows[0]
total_items = get_attack().total_affected_items
assert expected_total_items == total_items
def test_check_total_items_phase(mock_wdb, phase_name):
"""Test the number of returned items when filtering by phase."""
# load test database and make the query
cur = get_fake_database_data('schema_mitre_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT attack_id) FROM has_phase WHERE "
f"(phase_name='{phase_name}' COLLATE NOCASE)")
rows = cur.fetchone()
expected_total_items = rows[0]
total_items = get_attack(phase_name=phase_name).total_affected_items
assert expected_total_items == total_items
def test_check_total_items_searched_attack(mock_wdb, search):
"""Test the number of returned items when filtering by search."""
# load test database and make the query
cur = get_fake_database_data('schema_mitre_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT id) FROM"
f" attack WHERE json LIKE '%{search}%'")
rows = cur.fetchone()
expected_total_items = rows[0]
total_items = get_attack(search={'value': search, 'negation': 0}).total_affected_items
assert expected_total_items == total_items
def test_check_total_items_multiple_filters(mock_wdb, platform_name, phase_name):
"""Test the number of returned items when filtering by phase and platform.""" # noqa: E501
# load test database and make the query
cur = get_fake_database_data('schema_mitre_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT has_platform.attack_id) FROM "
"has_platform LEFT JOIN has_phase ON has_platform.attack_id = "
f"has_phase.attack_id WHERE (platform_name='{platform_name}' "
f" COLLATE NOCASE) AND (phase_name='{phase_name}' COLLATE NOCASE)")
rows = cur.fetchone()
expected_total_items = rows[0]
total_items = get_attack(platform_name=platform_name, phase_name=phase_name).total_affected_items
assert expected_total_items == total_items
def test_get_task_status_agent_id(mock_task_db, agent_id):
"""Check system's tasks (agent_id)
Parameters
----------
agent_id : str
Specific agent id
"""
cur = get_fake_database_data('schema_tasks_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT task_id) FROM tasks WHERE "
f"(agent_id='{int(agent_id)}' COLLATE NOCASE)")
rows = cur.fetchone()
expected_total_items = rows[0]
filters = {'agent_id': agent_id}
result = task.get_task_status(filters=filters)
assert result.total_affected_items == expected_total_items
def test_get_task_status_module(mock_task_db, module):
"""Check system's tasks (module)
Parameters
----------
module : str
Search tasks with a specific module
"""
cur = get_fake_database_data('schema_tasks_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT task_id) FROM tasks WHERE "
f"(module='{module}' COLLATE NOCASE)")
rows = cur.fetchone()
expected_total_items = rows[0]
filters = {'module': module}
result = task.get_task_status(filters=filters)
assert result.total_affected_items == expected_total_items
def test_get_task_status_command(mock_task_db, command):
"""Check system's tasks (command)
Parameters
----------
command : list
Search for tasks with a specific command
"""
cur = get_fake_database_data('schema_tasks_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT task_id) FROM tasks WHERE "
f"(command='{command}' COLLATE NOCASE)")
rows = cur.fetchone()
expected_total_items = rows[0]
filters = {'command': command}
result = task.get_task_status(filters=filters)
assert result.total_affected_items == expected_total_items
def test_get_task_status_node(mock_task_db, node):
"""Check system's tasks (node)
Parameters
----------
node : str
Search for the tasks of a specific node
"""
cur = get_fake_database_data('schema_tasks_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT task_id) FROM tasks WHERE "
f"(node='{node}' COLLATE NOCASE)")
rows = cur.fetchone()
expected_total_items = rows[0]
filters = {'node': node}
result = task.get_task_status(filters=filters)
assert result.total_affected_items == expected_total_items
def test_get_task_status_status(mock_task_db, status):
"""Check system's tasks (status)
Parameters
----------
status : str
Status of tasks to be shown
"""
cur = get_fake_database_data('schema_tasks_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT task_id) FROM tasks WHERE "
f"(status='{status}' COLLATE NOCASE)")
rows = cur.fetchone()
expected_total_items = rows[0]
filters = {'status': status}
result = task.get_task_status(filters=filters)
assert result.total_affected_items == expected_total_items
def test_get_attack_filter_limit(mock_wdb, limit, select):
"""Test if data are retrieved properly from Mitre database."""
result = get_attack(limit=limit, select=select)
# Max 10 results returned if json is included
if not select or 'json' in select:
expected_limit = min(10, limit)
assert len(
result.affected_items) <= expected_limit, f"Max expected results was 10, but {result.affected_items} returned."
else:
# Assert all results are returned
cur = get_fake_database_data('schema_mitre_test.sql').cursor()
cur.execute("SELECT COUNT(DISTINCT id) FROM attack")
rows = cur.fetchone()
expected_limit = min(rows[0], limit)
assert len(result.affected_items) <= expected_limit, f"Expected number or results was {expected_limit}, but " \
f"{len(result.affected_items)} returned."
def mitre_db():
"""Get fake MITRE database cursor."""
core_mitre.get_mitre_items.cache_clear()
return get_fake_database_data('schema_mitre_test.sql').cursor()
