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()