def run_local(agent_id=None, all_agents=False):
if agent_id == "000" or all_agents:
try:
SYSCHECK_RESTART = "{0}/var/run/.syscheck_run".format(
common.ossec_path)
fp = open(SYSCHECK_RESTART, 'w')
fp.write('{0}\n'.format(SYSCHECK_RESTART))
fp.close()
ret_msg = "Restarting Syscheck/Rootcheck locally"
except:
raise WazuhException(1601, "locally")
if all_agents:
oq = OssecQueue(common.ARQUEUE)
ret_msg = oq.send_msg_to_agent(OssecQueue.HC_SK_RESTART)
oq.close()
else:
# Check if agent exists and it is active
agent_info = Agent(agent_id).get_basic_information()
if 'status' in agent_info:
agent_status = agent_info['status']
else:
agent_status = "N/A"
if agent_status.lower() != 'active':
raise WazuhException(1602,
'{0} - {1}'.format(agent_id, agent_status))
oq = OssecQueue(common.ARQUEUE)
ret_msg = oq.send_msg_to_agent(OssecQueue.HC_SK_RESTART, agent_id)
oq.close()
return ret_msg
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_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_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(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_connect(mock_set, mock_conn):
"""Tests OssecQueue._connect function works"""
OssecQueue('test_path')
with patch('wazuh.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_command(agent_id=None, command=None, arguments=[], custom=False):
"""
Run AR command.
:param agent_id: Run AR command in the agent.
:return: Message.
"""
if not command:
raise WazuhException(1650, "Command not specified")
if not agent_id:
raise WazuhException(1650, "Agent ID not specified")
commands = get_commands()
if not custom and command not in commands:
raise WazuhException(1650, "Command not available")
# Create message
msg_queue = command
if custom:
msg_queue = "!{}".format(command)
if arguments:
msg_queue += " " + " ".join(shell_escape(str(x)) for x in arguments)
else:
msg_queue += " - -"
# Send
if agent_id == "000" or agent_id == "all":
oq = OssecQueue(common.EXECQ)
ret_msg = oq.send_msg_to_agent(msg=msg_queue,
agent_id=agent_id,
msg_type=OssecQueue.AR_TYPE)
oq.close()
if agent_id != "000" or agent_id == "all":
if agent_id != "all":
# Check if agent exists and it is active
agent_info = Agent(agent_id).get_basic_information()
if agent_info['status'].lower() != 'active':
raise WazuhException(1651)
if agent_id == "all":
agent_id = None
oq = OssecQueue(common.ARQUEUE)
ret_msg = oq.send_msg_to_agent(msg=msg_queue,
agent_id=agent_id,
msg_type=OssecQueue.AR_TYPE)
oq.close()
return ret_msg
def restart_agents(agent_id=None, restart_all=False):
"""
Restarts an agent or all agents.
:param agent_id: Agent ID of the agent to restart.
:param restart_all: Restarts all agents.
:return: Message.
"""
if restart_all:
oq = OssecQueue(common.ARQUEUE)
ret_msg = oq.send_msg_to_agent(OssecQueue.RESTART_AGENTS)
oq.close()
return ret_msg
else:
return Agent(agent_id).restart()
def restart_agents(agent_id=None, restart_all=False):
"""
Restarts an agent or all agents.
:param agent_id: Agent ID of the agent to restart.
:param restart_all: Restarts all agents.
:return: Message.
"""
if restart_all:
oq = OssecQueue(OssecQueue.ARQUEUE)
ret_msg = oq.send_msg_to_agent(OssecQueue.RESTART_AGENTS)
oq.close()
return ret_msg
else:
return Agent(agent_id).restart()
def restart(self):
"""
Restarts the agent.
:return: Message generated by OSSEC.
"""
if self.id == "000":
raise WazuhException(1703)
else:
# Check if agent exists and it is active
agent_info = self.get_basic_information()
if self.status.lower() != 'active':
raise WazuhException(1707, '{0} - {1}'.format(self.id, self.status))
oq = OssecQueue(common.ARQUEUE)
ret_msg = oq.send_msg_to_agent(OssecQueue.RESTART_AGENTS, self.id)
oq.close()
return ret_msg
def restart(self):
"""
Restarts the agent.
:return: Message generated by OSSEC.
"""
if self.id == "000":
raise WazuhException(1703)
else:
# Check if agent exists and it is active
agent_info = self.get_basic_information()
if self.status.lower() != 'active':
raise WazuhException(1707, '{0} - {1}'.format(self.id, self.status))
oq = OssecQueue(OssecQueue.ARQUEUE)
ret_msg = oq.send_msg_to_agent(OssecQueue.RESTART_AGENTS, self.id)
oq.close()
return ret_msg
def run(agent_id=None, all_agents=False):
"""
Runs rootcheck and syscheck.
:param agent_id: Run rootcheck/syscheck in the agent.
:param all_agents: Run rootcheck/syscheck in all agents.
:return: Message.
"""
if agent_id == "000" or all_agents:
try:
SYSCHECK_RESTART = "{0}/var/run/.syscheck_run".format(common.ossec_path)
fp = open(SYSCHECK_RESTART, 'w')
fp.write('{0}\n'.format(SYSCHECK_RESTART))
fp.close()
ret_msg = "Restarting Syscheck/Rootcheck locally"
except:
raise WazuhException(1601, "locally")
if all_agents:
oq = OssecQueue(OssecQueue.ARQUEUE)
ret_msg = oq.send_msg_to_agent(OssecQueue.HC_SK_RESTART)
oq.close()
else:
# Check if agent exists and it is active
agent_info = Agent(agent_id).get_basic_information()
if agent_info['status'].lower() != 'active':
raise WazuhException(1602, '{0} - {1}'.format(agent_id, agent_info['status']))
oq = OssecQueue(OssecQueue.ARQUEUE)
ret_msg = oq.send_msg_to_agent(OssecQueue.HC_SK_RESTART, agent_id)
oq.close()
return ret_msg
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 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()