diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index 2f5245d..1c17543 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -634,6 +634,24 @@ def install_package(self, host, package_name, become=None, ignore_errors=False): return result + def stop_ungracefully_process(self, host, process, become=None): + """Stop ungracefully a process. + + Args: + host (str): Hostname. + process (str): Name of the process to be stopped. + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + """ + try: + if self.is_windows(host): + self.run_command(host, f"taskkill /F /IM {process}", become=become) + else: + self.run_command(host, f"pkill -9 {process}", become=become) + + except Exception as e: + raise Exception(f"An error occurred while stopping the process. {e}") + def uninstall_package(self, host, package_name, become=None, ignore_errors=False): """Uninstall a package on a host. diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index 9346480..51a6b6e 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -3,6 +3,7 @@ # This program is free software; you can redistribute it and/or modify it under the terms of GPLv2 import os +import re import yaml from multiprocessing.pool import ThreadPool @@ -630,8 +631,8 @@ def restart_agent(self, host, method='service'): host (str): Hostname method (str): Method to restart agent """ - self.logger.debug(f'Restarting agent {host} via {method}') - service_name = WAZUH_ANGENT_WINDOWS_SERVICE_NAME if self.is_windows(host) else 'wazuh-agent' + self.logger.debug(f'Restarting agent {host}') + service_name = WAZUH_AGENT_WINDOWS_SERVICE_NAME if self.is_windows(host) else 'wazuh-agent' if self.is_agent(host): if method == 'service': self.control_service(host, service_name, 'restarted') @@ -677,7 +678,7 @@ def restart_managers(self, manager_list, parallel=True): """Restart managers Args: - manager_list (list): Managers list + manager_list (list): Managers list. parallel (bool, optional): Parallel execution. Defaults to True. """ self.logger.info(f"Restarting managers: {manager_list}") @@ -688,34 +689,42 @@ def restart_managers(self, manager_list, parallel=True): self.restart_manager(manager) self.logger.info(f"Managers restarted successfully: {manager_list}") - def stop_agent(self, host): - """Stop agent + def stop_agent(self, host, gracefully=True): + """Stop agent. Args: - host (str): Hostname + host (str): Hostname. + gracefully (bool): Stop gracefully. Defaults to True. """ - self.logger.debug(f'Stopping agent {host}') + type = 'gracefully' if gracefully else 'ungracefully' + self.logger.debug(f'Stopping {type} {host}') service_name = WAZUH_AGENT_WINDOWS_SERVICE_NAME if self.is_windows(host) else 'wazuh-agent' if self.is_agent(host): - self.control_service(host, service_name, 'stopped') - self.logger.debug(f"Agent {host} stopped successfully") + if not gracefully: + process = 'wazuh-agent.exe' if self.is_windows(host) else 'wazuh-agent' + self.stop_ungracefully_process(host, process) + + else: + self.control_service(host, service_name, 'stopped') + self.logger.info(f"Agent {host} stopped {type}") + else: raise ValueError(f"Host {host} is not an agent") - def stop_agents(self, agent_list=None, parallel=True): - """Stop agents + def stop_agents(self, agent_list=None, parallel=True, gracefully=True): + """Stop agents. Args: - agent_list(list, optional): Agents list. Defaults to None + agent_list (list, optional): Agents list. Defaults to None. parallel (bool, optional): Parallel execution. Defaults to True. + gracefully (bool): Stop gracefully. Defaults to True. """ - self.logger.info(f"Stopping agents: {agent_list}") if parallel: - self.pool.starmap(self.stop_agent, agent_list) + self.pool.starmap(self.stop_agent, [(agent, gracefully) for agent in agent_list]) else: for agent in agent_list: - self.restart_agent(agent) - self.logger.info(f"Agents stopped successfully: {agent_list}") + self.stop_agent(agent, gracefully) + self.logger.info(f'Agents stopped successfully: {agent_list}') def stop_manager(self, host): """Stop manager @@ -847,11 +856,11 @@ def stop_environment(self, parallel=True): self.pool.starmap(self.stop_agent, agent_list) else: self.logger.info(message='Stopping environment: Managers') - for manager in self.get_managers(): + for manager in manager_list: self.stop_manager(manager) self.logger.info(message='Stopping environment: Agents') - for agent in self.get_agents(): + for agent in agent_list: self.stop_agent(agent) self.logger.info('Stopping environment') @@ -874,11 +883,11 @@ def start_environment(self, parallel=True): self.pool.starmap(self.start_agent, agent_list) else: self.logger.info(message='Starting environment: Managers') - for manager in self.get_managers(): + for manager in manager_list: self.start_manager(manager) self.logger.info(message='Starting environment: Agents') - for agent in self.get_agents(): + for agent in agent_list: self.start_agent(agent) self.logger.info('Environment started successfully')