def test_WazuhSocketJSON_receive(mock_loads, mock_receive, mock_conn, raw):
"""Tests WazuhSocketJSON.receive function works"""
queue = WazuhSocketJSON('test_path')
response = queue.receive(raw=raw)
if raw:
assert isinstance(response, dict)
else:
assert isinstance(response, str)
mock_conn.assert_called_once_with('test_path')
def test_WazuhSocketJSON_receive_ko(mock_loads, mock_receive, mock_conn):
"""Tests WazuhSocketJSON.receive function works"""
queue = WazuhSocketJSON('test_path')
with pytest.raises(WazuhException, match=".* 10000 .*"):
queue.receive()
mock_conn.assert_called_once_with('test_path')
def test_WazuhSocketJSON_send(mock_send, mock_conn):
"""Tests WazuhSocketJSON.send function works"""
queue = WazuhSocketJSON('test_path')
response = queue.send('test_msg')
assert isinstance(response, MagicMock)
mock_conn.assert_called_once_with('test_path')
def send_logtest_msg(command: str = None, parameters: dict = None):
"""Connect and send a message to the logtest socket.
Parameters
----------
command: str
Command to send to the logtest socket.
parameters : dict
Dict of parameters that will be sent to the logtest socket.
Returns
-------
dict
Response from the logtest socket.
"""
full_message = create_wazuh_socket_message(origin={
'name': 'Logtest',
'module': 'api/framework'
},
command=command,
parameters=parameters)
logtest_socket = WazuhSocketJSON(LOGTEST_SOCKET)
logtest_socket.send(full_message)
response = logtest_socket.receive(raw=True)
logtest_socket.close()
return response
def send_logtest_msg(command: str = None, parameters: dict = None):
"""Connect and send a message to the logtest socket.
Parameters
----------
command: str
Command to send to the logtest socket.
parameters : dict
Dict of parameters that will be sent to the logtest socket.
Returns
-------
dict
Response from the logtest socket.
"""
full_message = create_wazuh_socket_message(origin={
'name': 'Logtest',
'module': origin_module.get()
},
command=command,
parameters=parameters)
logtest_socket = WazuhSocketJSON(LOGTEST_SOCKET)
logtest_socket.send(full_message)
response = logtest_socket.receive(raw=True)
logtest_socket.close()
try:
response['data']['output']['timestamp'] = datetime.strptime(
response['data']['output']['timestamp'],
"%Y-%m-%dT%H:%M:%S.%f+0000").strftime(decimals_date_format)
except KeyError:
pass
return response
def _remove_authd(self, purge=False):
"""Deletes the agent.
:param purge: Delete definitely from key store.
:return: Message.
"""
msg = {
"function": "remove",
"arguments": {
"id": str(self.id).zfill(3),
"purge": purge
}
}
authd_socket = WazuhSocketJSON(common.AUTHD_SOCKET)
authd_socket.send(msg)
data = authd_socket.receive()
authd_socket.close()
return data
def test_WazuhSocketJSON__init__(mock_conn):
"""Tests WazuhSocketJSON.__init__ function works"""
WazuhSocketJSON('test_path')
mock_conn.assert_called_once_with()
def _add_authd(self, name, ip, id=None, key=None, force=None):
"""Add an agent to Wazuh using authd.
2 uses:
- name and ip [force]: Add an agent like manage_agents (generate id and key).
- name, ip, id, key [force]: Insert an agent with an existing id and key.
Parameters
----------
name : str
Name of the new agent.
ip : str
IP of the new agent. It can be an IP, IP/NET or ANY.
id : str
ID of the new agent.
key : str
Key of the new agent.
force : dict
Remove old agents with same name or IP if conditions are met.
Raises
------
WazuhError(1705)
If there is an agent with the same name
WazuhError(1706)
If there is an agent with the same IP or the IP is invalid.
WazuhError(1708)
If there is an agent with the same ID.
WazuhError(1709)
If the key size is too short.
Returns
-------
Agent ID.
"""
# Check arguments
if id:
id = id.zfill(3)
if key and len(key) < 64:
raise WazuhError(1709)
msg = ""
if name and ip:
msg = {"function": "add", "arguments": {"name": name, "ip": ip}}
if force is not None:
# This force field must always be present
force.update({"key_mismatch": True})
msg["arguments"]["force"] = force
if id and key:
msg["arguments"].update({"id": id, "key": key})
try:
authd_socket = WazuhSocketJSON(common.AUTHD_SOCKET)
authd_socket.send(msg)
data = authd_socket.receive()
authd_socket.close()
except WazuhException as e:
if e.code == 9008:
raise WazuhError(1705, extra_message=name)
elif e.code == 9007:
raise WazuhError(1706, extra_message=ip)
elif e.code == 9012:
raise WazuhError(1708, extra_message=id)
raise e
self.id = data['id']
self.internal_key = data['key']
self.key = self.compute_key()