def test_remove_agents_database(send_mock, connect_mock, content):
"""
Tests delete_agents_db method handle exceptions properly
"""
def recv_mock(size_to_receive):
return bytes(len(content)) if size_to_receive == 4 else content
with patch('socket.socket.recv', side_effect=recv_mock):
mywdb = WazuhDBConnection()
received = mywdb.delete_agents_db(['001', '002'])
assert (isinstance(received, dict))
assert ("agents" in received)
def test_failed_send_private(send_mock, connect_mock):
"""
Tests an exception is properly raised when it's not possible to send a msg to the wdb socket
"""
def recv_mock(size_to_receive):
error_string = b'err {"agents": {"001": "Error"}}'
return bytes(len(error_string)) if size_to_receive == 4 else error_string
with patch('socket.socket.recv', side_effect=recv_mock):
mywdb = WazuhDBConnection()
with pytest.raises(exception.WazuhException, match=".* 2003 .*"):
mywdb._send('test_msg')
def test_wrong_character_encodings_wdb(send_mock, connect_mock):
"""
Tests receiving a text with a bad character encoding from wazuh db
"""
def recv_mock(size_to_receive):
bad_string = b' {"bad": "\x96bad"}'
return bytes(len(bad_string)) if size_to_receive == 4 else bad_string
with patch('socket.socket.recv', side_effect=recv_mock):
mywdb = WazuhDBConnection()
received = mywdb._send("test")
assert received == {"bad": "bad"}
def test_null_values_are_removed(send_mock, connect_mock):
"""
Tests '(null)' values are removed from the resulting dictionary
"""
def recv_mock(size_to_receive):
nulls_string = b' {"a": "a", "b": "(null)", "c": [1, 2, 3], "d": {"e": "(null)"}}'
return bytes(len(nulls_string)) if size_to_receive == 4 else nulls_string
with patch('socket.socket.recv', side_effect=recv_mock):
mywdb = WazuhDBConnection()
received = mywdb._send("test")
assert received == {"a": "a", "c": [1, 2, 3], "d": {}}
def test_failed_connection():
"""
Tests an exception is properly raised when it's not possible to connect to wdb
"""
# tests the socket path doesn't exists
with patch('wazuh.common.wdb_socket_path', '/this/path/doesnt/exist'):
with pytest.raises(exception.WazuhException, match=".* 2005 .*"):
WazuhDBConnection()
# tests an exception is properly raised when a connection error is raised
with patch('socket.socket') as socket_patch:
with pytest.raises(exception.WazuhException, match=".* 2005 .*"):
socket_patch.return_value.connect.side_effect = ConnectionError
WazuhDBConnection()
def clear(agent_id=None, all_agents=False):
"""
Clears the database.
:param agent_id: For an agent.
:param all_agents: For all agents.
:return: Message.
"""
agents = [agent_id] if not all_agents else map(
itemgetter('id'),
Agent.get_agents_overview(select={'fields': ['id']})['items'])
wdb_conn = WazuhDBConnection()
for agent in agents:
Agent(agent).get_basic_information() # check if the agent exists
wdb_conn.execute("agent {} sql delete from fim_entry".format(agent),
delete=True)
# update key fields which contains keys to value 000
wdb_conn.execute(
"agent {} sql update metadata set value = '000' where key like 'fim_db%'"
.format(agent),
update=True)
wdb_conn.execute(
"agent {} sql update metadata set value = '000' where key = 'syscheck-db-completed'"
.format(agent),
update=True)
return "Syscheck database deleted"
def __init__(self, agent_id, offset, limit, sort, search, select, query, count,
get_data, default_sort_field='policy_id', filters={}, fields=fields_translation_sca,
default_query=default_query_sca, count_field='policy_id'):
self.agent_id = agent_id
self._default_query_str = default_query
self.count_field = count_field
Agent(agent_id).get_basic_information() # check if the agent exists
WazuhDBQuery.__init__(self, offset=offset, limit=limit, table='sca_policy', sort=sort,
search=search, select=select, fields=fields, default_sort_field=default_sort_field,
default_sort_order='DESC', filters=filters, query=query, db_path=None,
min_select_fields=set(), count=count, get_data=get_data,
date_fields={'end_scan', 'start_scan'})
self.conn = WazuhDBConnection()
def __init__(self,
agent_id,
offset,
limit,
sort,
search,
select,
query,
count,
get_data,
default_sort_field='policy_id',
filters={},
fields=fields_translation_ca,
default_query=default_query_ca,
count_field='policy_id'):
self.agent_id = agent_id
self._default_query_str = default_query
self.count_field = count_field
Agent(agent_id).get_basic_information() # check if the agent exists
db_path = glob('{0}/{1}.db'.format(common.wdb_path, agent_id))
if not db_path:
raise WazuhException(1600)
WazuhDBQuery.__init__(self,
offset=offset,
limit=limit,
table='configuration_assessment_policy',
sort=sort,
search=search,
select=select,
fields=fields,
default_sort_field=default_sort_field,
default_sort_order='DESC',
filters=filters,
query=query,
db_path=db_path[0],
min_select_fields=set(),
count=count,
get_data=get_data,
date_fields={'end_scan', 'start_scan'})
self.conn = WazuhDBConnection()
def test_failed_execute(send_mock, connect_mock, error_query, error_type, expected_exception, delete, update):
mywdb = WazuhDBConnection()
if not error_type:
with pytest.raises(exception.WazuhException, match=f'.* {expected_exception} .*'):
mywdb.execute(error_query, delete=delete, update=update)
else:
with patch("wazuh.wdb.WazuhDBConnection._send", return_value=[{'total': 5}]):
with patch("wazuh.wdb.range", side_effect=error_type):
with pytest.raises(exception.WazuhException, match=f'.* {expected_exception} .*'):
mywdb.execute(error_query, delete=delete, update=update)
def connect_to_db(self):
return WazuhDBConnection()
def remove_bulk_agents(agent_ids_list: KeysView, logger):
"""
Removes files created by agents in worker nodes. This function doesn't remove agents from client.keys since the
client.keys file is overwritten by the master node.
:param agent_ids_list: List of agents ids to remove.
:param logger: Logger to use
:return: None.
"""
def remove_agent_file_type(agent_files: List[str]):
"""
Removes files if they exist
:param agent_files: Path regexes of the files to remove
:return: None
"""
for filetype in agent_files:
filetype_glob = filetype.format(ossec_path=common.ossec_path,
id='*',
name='*',
ip='*')
filetype_agent = {
filetype.format(ossec_path=common.ossec_path,
id=a['id'],
name=a['name'],
ip=a['ip'])
for a in agent_info
}
for agent_file in set(
glob.iglob(filetype_glob)) & filetype_agent:
logger.debug2("Removing {}".format(agent_file))
if os.path.isdir(agent_file):
shutil.rmtree(agent_file)
else:
os.remove(agent_file)
if not agent_ids_list:
return # the function doesn't make sense if there is no agents to remove
logger.info("Removing files from {} agents".format(
len(agent_ids_list)))
logger.debug("Agents to remove: {}".format(', '.join(agent_ids_list)))
# the agents must be removed in groups of 997: 999 is the limit of SQL variables per query. Limit and offset are
# always included in the SQL query, so that leaves 997 variables as limit.
for agents_ids_sublist in itertools.zip_longest(*itertools.repeat(
iter(agent_ids_list), 997),
fillvalue='0'):
agents_ids_sublist = list(
filter(lambda x: x != '0', agents_ids_sublist))
# Get info from DB
agent_info = Agent.get_agents_overview(
q=",".join(["id={}".format(i) for i in agents_ids_sublist]),
select={'fields': ['ip', 'id', 'name']},
limit=None)['items']
logger.debug2("Removing files from agents {}".format(
', '.join(agents_ids_sublist)))
files_to_remove = [
'{ossec_path}/queue/agent-info/{name}-{ip}',
'{ossec_path}/queue/rootcheck/({name}) {ip}->rootcheck',
'{ossec_path}/queue/diff/{name}',
'{ossec_path}/queue/agent-groups/{id}',
'{ossec_path}/queue/rids/{id}',
'{ossec_path}/var/db/agents/{name}-{id}.db'
]
remove_agent_file_type(files_to_remove)
logger.debug2("Removing agent group assigments from database")
# remove agent from groups
db_global = glob.glob(common.database_path_global)
if not db_global:
raise WazuhException(1600)
conn = Connection(db_global[0])
agent_ids_db = {
'id_agent{}'.format(i): int(i)
for i in agents_ids_sublist
}
conn.execute(
'delete from belongs where {}'.format(' or '.join([
'id_agent = :{}'.format(i) for i in agent_ids_db.keys()
])), agent_ids_db)
conn.commit()
# Tell wazuhbd to delete agent database
wdb_conn = WazuhDBConnection()
wdb_conn.delete_agents_db(agents_ids_sublist)
logger.info("Agent files removed")
def test_execute(send_mock, socket_send_mock, connect_mock):
mywdb = WazuhDBConnection()
mywdb.execute('agent 000 sql delete from test', delete=True)
mywdb.execute(
"agent 000 sql update test set value = 'test' where key = 'test'",
update=True)
with patch("wazuh.wdb.WazuhDBConnection._send",
return_value=[{
'total': 5
}]):
mywdb.execute("agent 000 sql select test from test offset 1 limit 1")
mywdb.execute("agent 000 sql select test from test offset 1 limit 1",
count=True)
mywdb.execute("agent 000 sql select test from test offset 1 count")
def test_query_lower_private(send_mock, connect_mock):
mywdb = WazuhDBConnection()
with pytest.raises(exception.WazuhException, match=".* 2004 .*"):
mywdb.execute("Agent sql select 'test'")
def test_query_input_validation_private(send_mock, connect_mock, error_query):
mywdb = WazuhDBConnection()
with pytest.raises(exception.WazuhException, match=".* 2004 .*"):
mywdb.execute(error_query)
class WazuhDBQueryPM(WazuhDBQuery):
def __init__(self,
agent_id,
offset,
limit,
sort,
search,
select,
query,
count,
get_data,
default_sort_field='policy_id',
filters={},
fields=fields_translation_ca,
default_query=default_query_ca,
count_field='policy_id'):
self.agent_id = agent_id
self._default_query_str = default_query
self.count_field = count_field
Agent(agent_id).get_basic_information() # check if the agent exists
db_path = glob('{0}/{1}.db'.format(common.wdb_path, agent_id))
if not db_path:
raise WazuhException(1600)
WazuhDBQuery.__init__(self,
offset=offset,
limit=limit,
table='configuration_assessment_policy',
sort=sort,
search=search,
select=select,
fields=fields,
default_sort_field=default_sort_field,
default_sort_order='DESC',
filters=filters,
query=query,
db_path=db_path[0],
min_select_fields=set(),
count=count,
get_data=get_data,
date_fields={'end_scan', 'start_scan'})
self.conn = WazuhDBConnection()
def _default_query(self):
return self._default_query_str
def _substitute_params(self):
for k, v in self.request.items():
self.query = self.query.replace(f':{k}', f"'{v}'")
def _get_total_items(self):
self._substitute_params()
total_items = self.conn.execute(
f'agent {self.agent_id} sql ' +
self.query.format(self._default_count_query()))
self.total_items = total_items if isinstance(
total_items, int) else total_items[0][self._default_count_query()]
def _default_count_query(self):
return f"COUNT(DISTINCT {self.count_field})"
def _get_data(self):
self._substitute_params()
self._data = self.conn.execute(
f'agent {self.agent_id} sql ' + self.query.format(','.join(
map(lambda x: self.fields[x], self.select['fields']
| self.min_select_fields))))
def _format_data_into_dictionary(self):
return {"totalItems": self.total_items, "items": self._data}
def _add_limit_to_query(self):
if self.limit:
if self.limit > common.maximum_database_limit:
raise WazuhException(1405, str(self.limit))
self.query += f' LIMIT {self.limit} OFFSET {self.offset}'
elif self.limit == 0: # 0 is not a valid limit
raise WazuhException(1406)
def run(self):
self._add_select_to_query()
self._add_filters_to_query()
self._add_search_to_query()
if self.count:
self._get_total_items()
self._add_sort_to_query()
self._add_limit_to_query()
if self.data:
self._get_data()
return self._format_data_into_dictionary()