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")
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()
