def __init__(self, inventory_path, messages_path, tmp_path, time_step=0.5):
"""Create a new instance to monitor any given file in any specified host.
Parameters
----------
inventory_path : str
Path to the hosts's inventory file.
messages_path : str
Path to the file where the callbacks, paths and hosts to be monitored are specified.
tmp_path : str
Path to the temporal files.
time_step : float, optional
Fraction of time to wait in every get. Default `0.5`.
"""
self.host_manager = HostManager(inventory_path=inventory_path)
self._queue = Manager().Queue()
self._result = defaultdict(list)
self._time_step = time_step
self._file_monitors = list()
self._monitored_files = set()
self._file_content_collectors = list()
self._tmp_path = tmp_path
try:
os.mkdir(self._tmp_path)
except OSError:
pass
with open(messages_path, 'r') as f:
self.test_cases = yaml.safe_load(f)
def test_agent_key_polling(inventory_path):
"""Check that the agent key polling cycle works correctly. To do this, we use the messages and the hosts defined
in data/messages.yml and the hosts inventory.
Parameters
----------
inventory_path : str
Path to the Ansible hosts inventory
"""
actual_path = os.path.dirname(os.path.abspath(__file__))
host_manager = HostManager(inventory_path=inventory_path)
configure_environment(host_manager)
host_monitor = HostMonitor(inventory_path=inventory_path,
messages_path=os.path.join(
actual_path, 'data/messages.yml'),
tmp_path=os.path.join(actual_path, 'tmp'))
host_monitor.run()
test_hosts = ['wazuh-master', 'wazuh-worker1', 'wazuh-worker2']
inventory_path = os.path.join(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
'provisioning', 'agentless_cluster', 'inventory.yml')
default_api_conf = os.path.join(os.path.dirname(os.path.abspath(__file__)),
'api_configurations', 'default.yaml')
# Testing credentials
test_user = None
test_passw = None
test_user_id = None
test_role_id = None
test_rule_id = None
test_policy_id = None
host_manager = HostManager(inventory_path)
@pytest.fixture(scope='module')
def set_role_to_user():
"""Create a relation between the testing user and role/policy."""
token = host_manager.get_api_token(test_hosts[0])
response = host_manager.make_api_call(
test_hosts[0],
method='POST',
endpoint=
f'/security/users/{test_user_id}/roles?role_ids={test_role_id}',
token=token)
assert response[
'status'] == 200, f'Failed to set relation between user and role: {response}'
class HostMonitor:
"""This class has the capability to monitor remote host. This monitoring consists of reading the specified files to
check that the expected message arrives to them.
If the goals are achieved, no exceptions will be raised and therefore the test will end properly and without
failures.
In contrast, if one or more of the goals is not covered, a timeout exception will be raised with a generic or a
custom error message.
"""
def __init__(self, inventory_path, messages_path, tmp_path, time_step=0.5):
"""Create a new instance to monitor any given file in any specified host.
Parameters
----------
inventory_path : str
Path to the hosts's inventory file.
messages_path : str
Path to the file where the callbacks, paths and hosts to be monitored are specified.
tmp_path : str
Path to the temporal files.
time_step : float, optional
Fraction of time to wait in every get. Default `0.5`.
"""
self.host_manager = HostManager(inventory_path=inventory_path)
self._queue = Manager().Queue()
self._result = defaultdict(list)
self._time_step = time_step
self._file_monitors = list()
self._monitored_files = set()
self._file_content_collectors = list()
self._tmp_path = tmp_path
try:
os.mkdir(self._tmp_path)
except OSError:
pass
with open(messages_path, 'r') as f:
self.test_cases = yaml.safe_load(f)
def run(self):
"""This method creates and destroy the needed processes for the messages founded in messages_path.
It creates one file composer (process) for every file to be monitored in every host."""
for host, payload in self.test_cases.items():
self._monitored_files.update({case['path'] for case in payload})
if len(self._monitored_files) == 0:
raise AttributeError('There is no path to monitor. Exiting...')
for path in self._monitored_files:
output_path = f'{host}_{path.split("/")[-1]}.tmp'
self._file_content_collectors.append(self.file_composer(host=host, path=path, output_path=output_path))
logger.debug(f'Add new file composer process for {host} and path: {path}')
self._file_monitors.append(self._start(host=host, payload=payload, path=output_path))
logger.debug(f'Add new file monitor process for {host} and path: {path}')
while True:
if not any([handler.is_alive() for handler in self._file_monitors]):
for handler in self._file_monitors:
handler.join()
for file_collector in self._file_content_collectors:
file_collector.terminate()
file_collector.join()
self.clean_tmp_files()
break
time.sleep(self._time_step)
self.check_result()
@new_process
def file_composer(self, host, path, output_path):
"""Collects the file content of the specified path in the desired host and append it to the output_path file.
Simulates the behavior of tail -f and redirect the output to output_path.
Parameters
----------
host : str
Hostname.
path : str
Host file path to be collect.
output_path : str
Output path of the content collected from the remote host path.
"""
try:
truncate_file(os.path.join(self._tmp_path, output_path))
except FileNotFoundError:
pass
logger.debug(f'Starting file composer for {host} and path: {path}. '
f'Composite file in {os.path.join(self._tmp_path, output_path)}')
tmp_file = os.path.join(self._tmp_path, output_path)
while True:
with FileLock(tmp_file):
with open(tmp_file, "r+") as file:
content = self.host_manager.get_file_content(host, path).split('\n')
file_content = file.read().split('\n')
for new_line in content:
if new_line == '':
continue
if new_line not in file_content:
file.write(f'{new_line}\n')
time.sleep(self._time_step)
@new_process
def _start(self, host, payload, path, encoding=None, error_messages_per_host=None):
"""Start the file monitoring until the QueueMonitor returns an string or TimeoutError.
Parameters
----------
host : str
Hostname
payload : list of dict
Contains the message to be found and the timeout for it.
path : str
Path where it must search for the message.
encoding : str
Encoding of the file.
error_messages_per_host : dict
Dictionary with hostnames as keys and desired error messages as values
Returns
-------
instance of HostMonitor
"""
tailer = FileTailer(os.path.join(self._tmp_path, path), time_step=self._time_step)
try:
if encoding is not None:
tailer.encoding = encoding
tailer.start()
for case in payload:
logger.debug(f'Starting QueueMonitor for {host} and message: {case["regex"]}')
monitor = QueueMonitor(tailer.queue, time_step=self._time_step)
try:
self._queue.put({host: monitor.start(timeout=case['timeout'],
callback=callback_generator(case['regex'])
).result().strip('\n')})
except TimeoutError:
try:
self._queue.put({host: error_messages_per_host[host]})
except (KeyError, TypeError):
self._queue.put({
host: TimeoutError(f'Did not found the expected callback in {host}: {case["regex"]}')})
logger.debug(f'Finishing QueueMonitor for {host} and message: {case["regex"]}')
finally:
tailer.shutdown()
return self
def result(self):
"""Get the result of HostMonitor
Returns
-------
dict
Dict that contains the host as the key and a list of messages as the values
"""
return self._result
def check_result(self):
"""Check if a TimeoutError occurred."""
logger.debug(f'Checking results...')
while not self._queue.empty():
result = self._queue.get(block=True)
for host, msg in result.items():
if isinstance(msg, TimeoutError):
raise msg
logger.debug(f'Received from {host} the expected message: {msg}')
self._result[host].append(msg)
def clean_tmp_files(self):
"""Remove tmp files."""
logger.debug(f'Cleaning temporal files...')
for file in os.listdir(self._tmp_path):
os.remove(os.path.join(self._tmp_path, file))