def test_OssecQueue_protected_connect(mock_set, mock_conn):
"""Tests OssecQueue._connect function works"""
OssecQueue('test_path')
with patch('wazuh.core.ossec_queue.socket.socket.getsockopt', return_value=1):
OssecQueue('test_path')
mock_conn.assert_called_with('test_path')
mock_set.assert_called_once_with(1, 7, 6400)
def run(agent_list=None):
"""Run syscheck scan.
:param agent_list: Run syscheck in the agent.
:return: AffectedItemsWazuhResult.
"""
result = AffectedItemsWazuhResult(all_msg='Syscheck scan was restarted on returned agents',
some_msg='Syscheck scan was not restarted on some agents',
none_msg='No syscheck scan was restarted')
for agent_id in agent_list:
try:
agent_info = Agent(agent_id).get_basic_information()
agent_status = agent_info.get('status', 'N/A')
if agent_status.lower() != 'active':
result.add_failed_item(
id_=agent_id, error=WazuhError(1601, extra_message='Status - {}'.format(agent_status)))
else:
oq = OssecQueue(common.ARQUEUE)
oq.send_msg_to_agent(OssecQueue.HC_SK_RESTART, agent_id)
result.affected_items.append(agent_id)
oq.close()
except WazuhError as e:
result.add_failed_item(id_=agent_id, error=e)
result.affected_items = sorted(result.affected_items, key=int)
result.total_affected_items = len(result.affected_items)
return result
def run_command(agent_list=None, command=None, arguments=None, custom=False):
"""Run AR command in a specific agent
:param agent_list: Run AR command in the agent.
:param command: Command running in the agent. If this value starts by !, then it refers to a script name instead of
a command name
:param custom: Whether the specified command is a custom command or not
:param arguments: Command arguments
:return: AffectedItemsWazuhResult.
"""
msg_queue = active_response.create_message(command=command,
arguments=arguments,
custom=custom)
oq = OssecQueue(common.ARQUEUE)
result = AffectedItemsWazuhResult(
none_msg='Could not send command to any agent',
some_msg='Could not send command to some agents',
all_msg='Command sent to all agents')
for agent_id in agent_list:
try:
active_response.send_command(msg_queue, oq, agent_id)
result.affected_items.append(agent_id)
result.total_affected_items += 1
except WazuhException as e:
result.add_failed_item(id_=agent_id, error=e)
oq.close()
return result
def test_OssecQueue_close(mock_close, mock_conn):
"""Tests OssecQueue.close function works"""
queue = OssecQueue('test_path')
queue.close()
mock_conn.assert_called_once_with('test_path')
mock_close.assert_called_once_with()
def test_OssecQueue_protected_send_ko(mock_send, mock_conn):
"""Tests OssecQueue._send function exceptions works"""
queue = OssecQueue('test_path')
with pytest.raises(WazuhException, match=".* 1011 .*"):
queue._send('msg')
mock_conn.assert_called_with('test_path')
def test_OssecQueue_send_msg_to_agent_ko(mock_send, mock_conn, msg, agent_id, msg_type, expected_exception):
"""Tests OssecQueue.send_msg_to_agent function exception works"""
queue = OssecQueue('test_path')
with pytest.raises(WazuhException, match=f'.* {expected_exception} .*'):
queue.send_msg_to_agent(msg, agent_id, msg_type)
mock_conn.assert_called_once_with('test_path')
def test_OssecQueue_send_msg_to_agent(mock_send, mock_conn, msg, agent_id, msg_type):
"""Tests OssecQueue.send_msg_to_agent function works"""
queue = OssecQueue('test_path')
response = queue.send_msg_to_agent(msg, agent_id, msg_type)
assert isinstance(response, str)
mock_conn.assert_called_once_with('test_path')
def test_OssecQueue_protected_send(mock_conn, send_response, error):
"""Tests OssecQueue._send function works"""
queue = OssecQueue('test_path')
with patch('socket.socket.send', return_value=send_response):
if error:
with pytest.raises(WazuhException, match=".* 1011 .*"):
queue._send('msg')
else:
queue._send('msg')
mock_conn.assert_called_with('test_path')
def run_command(agent_list: list = None,
command: str = '',
arguments: list = None,
custom: bool = False,
alert: dict = None) -> AffectedItemsWazuhResult:
"""Run AR command in a specific agent.
Parameters
----------
agent_list : list
Agents list that will run the AR command.
command : str
Command running in the agents. If this value starts with !, then it refers to a script name instead of a
command name.
custom : bool
Whether the specified command is a custom command or not.
arguments : list
Command arguments.
alert : dict
Alert information depending on the AR executed.
Returns
-------
AffectedItemsWazuhResult.
"""
result = AffectedItemsWazuhResult(
all_msg='AR command was sent to all agents',
some_msg='AR command was not sent to some agents',
none_msg='AR command was not sent to any agent')
if agent_list:
oq = OssecQueue(common.ARQUEUE)
system_agents = get_agents_info()
for agent_id in agent_list:
try:
if agent_id not in system_agents:
raise WazuhResourceNotFound(1701)
if agent_id == "000":
raise WazuhError(1703)
active_response.send_ar_message(agent_id, oq, command,
arguments, custom, alert)
result.affected_items.append(agent_id)
result.total_affected_items += 1
except WazuhException as e:
result.add_failed_item(id_=agent_id, error=e)
oq.close()
return result
def test_OssecQueue_protected_connect_ko(mock_conn):
"""Tests OssecQueue._connect function exceptions works"""
with pytest.raises(WazuhException, match=".* 1010 .*"):
OssecQueue('test_path')
def test_OssecQueue__init__(mock_conn):
"""Tests OssecQueue.__init__ function works"""
OssecQueue('test_path')
mock_conn.assert_called_once_with()