def send_restart_command(agent_id: str = '',
agent_version: str = '',
wq: WazuhQueue = None) -> str:
"""Send restart command to an agent.
Parameters
----------
agent_id : str
ID specifying the agent where the restart command will be sent to
agent_version : str
Agent version to compare with the required version. The format is vX.Y.Z.
wq : WazuhQueue
WazuhQueue used for the active response messages.
Returns
-------
str
Message generated by Wazuh.
"""
# If the Wazuh agent version is newer or equal to the AR legacy version,
# the message sent will have JSON format
if WazuhVersion(agent_version) >= WazuhVersion(common.AR_LEGACY_VERSION):
ret_msg = wq.send_msg_to_agent(WazuhQueue.RESTART_AGENTS_JSON,
agent_id)
else:
ret_msg = wq.send_msg_to_agent(WazuhQueue.RESTART_AGENTS, agent_id)
return ret_msg
def reconnect_agents(agent_list: Union[list, str] = None) -> AffectedItemsWazuhResult:
"""Force reconnect a list of agents.
Parameters
----------
agent_list : Union[list, str]
List of agent IDs. All possible values from 000 onwards. Default `*`
Returns
-------
AffectedItemsWazuhResult
"""
result = AffectedItemsWazuhResult(all_msg='Force reconnect command was sent to all agents',
some_msg='Force reconnect command was not sent to some agents',
none_msg='Force reconnect command was not sent to any agent'
)
system_agents = get_agents_info()
wq = WazuhQueue(common.ARQUEUE)
for agent_id in agent_list:
try:
if agent_id not in system_agents:
raise WazuhResourceNotFound(1701)
if agent_id == "000":
raise WazuhError(1703)
Agent(agent_id).reconnect(wq)
result.affected_items.append(agent_id)
except WazuhException as e:
result.add_failed_item(id_=agent_id, error=e)
wq.close()
result.total_affected_items = len(result.affected_items)
result.affected_items.sort(key=int)
return result
def send_ar_message(agent_id: str = '',
wq: WazuhQueue = None,
command: str = '',
arguments: list = None,
custom: bool = False,
alert: dict = None) -> None:
"""Send the active response message to the agent.
Parameters
----------
agent_id : str
ID specifying the agent where the msg_queue will be sent to.
wq : WazuhQueue
WazuhQueue used for the active response messages.
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.
Raises
------
WazuhError(1651)
If the agent with ID agent_id is not active.
"""
# Agent basic information
agent_info = Agent(agent_id).get_basic_information()
# Check if agent is active
if agent_info['status'].lower() != 'active':
raise WazuhError(1651,
extra_message='{0}'.format(agent_info['status']))
# Once we know the agent is active, store version
agent_version = agent_info['version']
# Check if AR is enabled
agent_conf = Agent(agent_id).getconfig('com', 'active-response',
agent_version)
if agent_conf['active-response']['disabled'] == 'yes':
raise WazuhError(1750)
# Create classic msg or JSON msg depending on the agent version
if WazuhVersion(agent_version) >= WazuhVersion(common.AR_LEGACY_VERSION):
msg_queue = create_json_message(command=command,
arguments=arguments,
alert=alert)
else:
msg_queue = create_message(command=command,
arguments=arguments,
custom=custom)
wq.send_msg_to_agent(msg=msg_queue,
agent_id=agent_id,
msg_type=WazuhQueue.AR_TYPE)
def test_WazuhQueue_close(mock_close, mock_conn):
"""Tests WazuhQueue.close function works"""
queue = WazuhQueue('test_path')
queue.close()
mock_conn.assert_called_once_with('test_path')
mock_close.assert_called_once_with()
def test_WazuhQueue_protected_send_ko(mock_send, mock_conn):
"""Tests WazuhQueue._send function exceptions works"""
queue = WazuhQueue('test_path')
with pytest.raises(WazuhException, match=".* 1011 .*"):
queue._send('msg')
mock_conn.assert_called_with('test_path')
def test_WazuhQueue_send_msg_to_agent(mock_send, mock_conn, msg, agent_id,
msg_type):
"""Tests WazuhQueue.send_msg_to_agent function works"""
queue = WazuhQueue('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_WazuhQueue_send_msg_to_agent_ko(mock_send, mock_conn, msg, agent_id,
msg_type, expected_exception):
"""Tests WazuhQueue.send_msg_to_agent function exception works"""
queue = WazuhQueue('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_WazuhQueue_protected_connect(mock_set, mock_conn):
"""Tests WazuhQueue._connect function works"""
WazuhQueue('test_path')
with patch('wazuh.core.wazuh_queue.socket.socket.getsockopt',
return_value=1):
WazuhQueue('test_path')
mock_conn.assert_called_with('test_path')
mock_set.assert_called_once_with(1, 7, 6400)
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:
wq = WazuhQueue(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, wq, 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)
result.affected_items.sort(key=int)
wq.close()
return result
def reconnect(self, wq: WazuhQueue) -> str:
"""Force reconnect to the manager.
Parameters
----------
wq : WazuhQueue
WazuhQueue used for the active response message.
Raises
------
WazuhError(1750)
If the agent has active response disabled.
WazuhError(1757)
If the agent to be reconnected is not active.
Returns
-------
str
Message generated by Wazuh.
"""
# Check if agent is active
self.get_basic_information()
if self.status.lower() != 'active':
raise WazuhError(1757)
# Send force reconnect message to the WazuhQueue
ret_msg = wq.send_msg_to_agent(WazuhQueue.HC_FORCE_RECONNECT, self.id)
return ret_msg
def test_WazuhQueue_send_msg_to_agent(mock_send, mock_conn, msg, agent_id,
msg_type):
"""Test WazuhQueue.send_msg_to_agent function.
Parameters
----------
msg : str
Message sent to the agent.
agent_id : str
String indicating the agent ID.
msg_type : str
String indicating the message type.
"""
queue = WazuhQueue('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 run(agent_list: Union[list, None] = None) -> AffectedItemsWazuhResult:
"""Run a rootcheck scan in the specified agents.
Parameters
----------
agent_list : Union[list, None]
List of the agents IDs to run the scan for.
Returns
-------
result : AffectedItemsWazuhResult
JSON containing the affected agents.
"""
result = AffectedItemsWazuhResult(
all_msg='Rootcheck scan was restarted on returned agents',
some_msg='Rootcheck scan was not restarted on some agents',
none_msg='No rootcheck scan was restarted')
system_agents = get_agents_info()
rbac_filters = get_rbac_filters(system_resources=system_agents,
permitted_resources=agent_list)
agent_list = set(agent_list)
not_found_agents = agent_list - system_agents
# Add non existent agents to failed_items
[
result.add_failed_item(id_=agent, error=WazuhResourceNotFound(1701))
for agent in not_found_agents
]
# Add non eligible agents to failed_items
non_eligible_agents = WazuhDBQueryAgents(limit=None,
select=["id", "status"],
query=f'status!=active',
**rbac_filters).run()['items']
[
result.add_failed_item(
id_=agent['id'],
error=WazuhError(1601,
extra_message=f'Status - {agent["status"]}'))
for agent in non_eligible_agents
]
wq = WazuhQueue(common.ARQUEUE)
eligible_agents = agent_list - not_found_agents - {
d['id']
for d in non_eligible_agents
}
for agent_id in eligible_agents:
try:
wq.send_msg_to_agent(WazuhQueue.HC_SK_RESTART, agent_id)
result.affected_items.append(agent_id)
except WazuhError as e:
result.add_failed_item(id_=agent_id, error=e)
wq.close()
result.affected_items.sort(key=int)
result.total_affected_items = len(result.affected_items)
return result
def test_WazuhQueue_send_msg_to_agent_ko(mock_send, mock_conn, msg, agent_id,
msg_type, expected_exception):
"""Test WazuhQueue.send_msg_to_agent function exceptions.
Parameters
----------
msg : str
Message sent to the agent.
agent_id : str
String indicating the agent ID.
msg_type : str
String indicating the message type.
expected_exception : int
Expected Wazuh exception.
"""
queue = WazuhQueue('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_WazuhQueue_protected_send(mock_conn, send_response, error):
"""Tests WazuhQueue._send function works"""
queue = WazuhQueue('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 test_WazuhQueue_protected_send(mock_conn, send_response, error):
"""Test WazuhQueue._send function.
Parameters
----------
send_response : int
Returned value of the socket send mocked function.
error : bool
Indicates whether a WazuhException will be raised or not.
"""
queue = WazuhQueue('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(agent_list=None):
"""Run a syscheck scan in the specified agents.
Parameters
----------
agent_list : str
List of the agents IDs to run the scan for.
Returns
-------
result : AffectedItemsWazuhResult
Confirmation/Error message.
"""
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:
wq = WazuhQueue(common.ARQUEUE)
wq.send_msg_to_agent(WazuhQueue.HC_SK_RESTART, agent_id)
result.affected_items.append(agent_id)
wq.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 restart_agents(agent_list: list = None) -> AffectedItemsWazuhResult:
"""Restart a list of agents.
Parameters
----------
agent_list : list
List of agents IDs.
Raises
------
WazuhError(1703)
If the agent to be restarted is 000.
WazuhError(1701)
If the agent to be restarted is not in the system.
WazuhError(1707)
If the agent to be restarted is not active.
Returns
-------
AffectedItemsWazuhResult
"""
result = AffectedItemsWazuhResult(
all_msg='Restart command was sent to all agents',
some_msg='Restart command was not sent to some agents',
none_msg='Restart command was not sent to any agent')
agent_list = set(agent_list)
# Add agent with ID 000 to failed_items
try:
agent_list.remove('000')
result.add_failed_item('000', WazuhError(1703))
except KeyError:
pass
if agent_list:
system_agents = get_agents_info()
rbac_filters = get_rbac_filters(system_resources=system_agents,
permitted_resources=list(agent_list))
agents_with_data = WazuhDBQueryAgents(
limit=None, select=["id", "status",
"version"], **rbac_filters).run()['items']
# Add non existent agents to failed_items
not_found_agents = agent_list - system_agents
[
result.add_failed_item(id_=agent,
error=WazuhResourceNotFound(1701))
for agent in not_found_agents
]
# Add non active agents to failed_items
non_active_agents = [
agent for agent in agents_with_data if agent['status'] != 'active'
]
[
result.add_failed_item(id_=agent['id'],
error=WazuhError(
1707,
extra_message=f'{agent["status"]}'))
for agent in non_active_agents
]
eligible_agents = [agent for agent in agents_with_data if agent not in non_active_agents] if non_active_agents \
else agents_with_data
wq = WazuhQueue(common.ARQUEUE)
for agent in eligible_agents:
try:
send_restart_command(agent['id'], agent['version'], wq)
result.affected_items.append(agent['id'])
except WazuhException as e:
result.add_failed_item(id_=agent['id'], error=e)
wq.close()
result.total_affected_items = len(result.affected_items)
result.affected_items.sort(key=int)
return result
def test_WazuhQueue__init__(mock_conn):
"""Tests WazuhQueue.__init__ function works"""
WazuhQueue('test_path')
mock_conn.assert_called_once_with()
def test_WazuhQueue_protected_connect_ko(mock_conn):
"""Tests WazuhQueue._connect function exceptions works"""
with pytest.raises(WazuhException, match=".* 1010 .*"):
WazuhQueue('test_path')
