From a3338d7aaab161e20f4c7e529c96692ddddabe82 Mon Sep 17 00:00:00 2001 From: Antonio Date: Wed, 28 Jun 2023 14:39:10 +0200 Subject: [PATCH 1/6] feature(#61): stopping agents ungracefully --- .../system/wazuh_handler.py | 45 +++++++++++++------ 1 file changed, 31 insertions(+), 14 deletions(-) diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index 59f7e02..ccb4755 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -4,6 +4,7 @@ import os import re +from functools import partial from multiprocessing.pool import ThreadPool from wazuh_qa_framework.generic_modules.logging.base_logger import BaseLogger @@ -568,33 +569,49 @@ 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): + def stop_agent(self, host, type='gracefully'): """Stop agent Args: host (str): Hostname + type (str): Type of stop, it can be gracefully or ungracefully. Defaults to gracefully """ - self.logger.debug(f'Stopping agent {host}') - service_name = WAZUH_ANGENT_WINDOWS_SERVICE_NAME if is_windows(host) else 'wazuh-agent' + self.logger.debug(f'Stopping agent {type} {host}') + service_name = WAZUH_ANGENT_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 type == 'gracefully': + if self.is_windows(host): + self.control_service(host, service_name, 'stopped') + self.logger.debug(f'Agent {host} stopped successfully') + else: + self.control_service(host, service_name, 'stopped', become=True) + self.logger.debug(f'Agent {host} stopped successfully') + + if type == 'ungracefully': + if self.is_windows(host): + self.run_command(host, 'taskkill /F /IM wazuh-agent.exe') + self.logger.debug(f'Agent {host} stopped successfully') + else: + self.run_command(host, "pkill -9 wazuh-agent", become=True) + self.logger.debug(f'Agent {host} stopped successfully') else: raise ValueError(f'Host {host} is not an agent') - def stop_agents(self, agent_list=None, parallel=True): + def stop_agents(self, agent_list=None, parallel=True, type='gracefully'): """Stop agents Args: agent_list(list, optional): Agents list. Defaults to None parallel (bool, optional): Parallel execution. Defaults to True. + type (str): Type of stop, it can be gracefully or ungracefully. Defaults to gracefully """ - self.logger.info(f'Stopping agents: {agent_list}') + self.logger.info(f'Stopping agents {type}: {agent_list}') if parallel: - self.pool.map(self.stop_agent, agent_list) + stop_agent_with_type = partial(self.stop_agent, type=type) + self.pool.map(stop_agent_with_type, agent_list) else: for agent in agent_list: - self.restart_agent(agent) + self.stop_agent(agent, type) self.logger.info(f'Agents stopped successfully: {agent_list}') def stop_manager(self, host): @@ -632,7 +649,7 @@ def start_agent(self, host): host (str): Hostname """ self.logger.debug(f'Starting agent {host}') - service_name = WAZUH_ANGENT_WINDOWS_SERVICE_NAME if is_windows(host) else 'wazuh-agent' + service_name = WAZUH_ANGENT_WINDOWS_SERVICE_NAME if self.is_windows(host) else 'wazuh-agent' if self.is_agent(host): self.control_service(host, service_name, 'started') self.logger.debug(f'Agent {host} started successfully') @@ -727,11 +744,11 @@ def stop_environment(self, parallel=True): self.pool.map(self.stop_agent, agent_list) else: self.logger.info(message='Stopping environment: Managers') - for manager in get_managers(): + for manager in manager_list: self.stop_manager(manager) self.logger.info(message='Stopping environment: Agents') - for agent in get_agents(): + for agent in manager_list: self.stop_agent(agent) self.logger.info('Stopping environment') @@ -754,11 +771,11 @@ def start_environment(self, parallel=True): self.pool.map(self.start_agent, agent_list) else: self.logger.info(message='Starting environment: Managers') - for manager in get_managers(): + for manager in manager_list: self.start_manager(manager) self.logger.info(message='Starting environment: Agents') - for agent in get_agents(): + for agent in agent_list : self.start_agent(agent) self.logger.info('Environment started successfully') From 62b6339d706ed6e0f3222021702651dd4a0723a7 Mon Sep 17 00:00:00 2001 From: Antonio Date: Wed, 28 Jun 2023 14:48:06 +0200 Subject: [PATCH 2/6] fix(#61): fixes after linter --- src/wazuh_qa_framework/system/wazuh_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index ccb4755..bedc189 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -775,7 +775,7 @@ def start_environment(self, parallel=True): self.start_manager(manager) self.logger.info(message='Starting environment: Agents') - for agent in agent_list : + for agent in agent_list: self.start_agent(agent) self.logger.info('Environment started successfully') From ab390ad2c27a3f205a50472e7c0c3a9dc0368f8c Mon Sep 17 00:00:00 2001 From: Antonio Date: Thu, 13 Jul 2023 13:29:10 +0200 Subject: [PATCH 3/6] fix(#61): changes after review --- .../global_variables/daemons.py | 2 +- src/wazuh_qa_framework/system/host_manager.py | 16 +++++ .../system/wazuh_handler.py | 65 +++++++++---------- 3 files changed, 46 insertions(+), 37 deletions(-) diff --git a/src/wazuh_qa_framework/global_variables/daemons.py b/src/wazuh_qa_framework/global_variables/daemons.py index fce1140..6ce2f7e 100644 --- a/src/wazuh_qa_framework/global_variables/daemons.py +++ b/src/wazuh_qa_framework/global_variables/daemons.py @@ -18,7 +18,7 @@ WAZUH_DB_DAEMON = 'wazuh-db' AGENT_DAEMON = 'wazuh-agentd' -WAZUH_ANGENT_WINDOWS_SERVICE_NAME = 'WazuhSvc' +WAZUH_AGENT_WINDOWS_SERVICE_NAME = 'WazuhSvc' WAZUH_AGENT_DAEMONS = [LOGCOLLECTOR_DAEMON, SYSCHECK_DAEMON, EXEC_DAEMON, MODULES_DAEMON, AGENT_DAEMON] diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index 668ff98..4711d0b 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -632,6 +632,22 @@ def install_package(self, host, package_name, become=None, ignore_errors=False): return result + def stop_ungracefully_process(self, host, process): + """Stop ungracefully a process. + + Args: + host (str): Hostname. + process (str): Name of the process to be stopped. + """ + try: + if self.is_windows(host): + self.run_command(host, f'taskkill /F /IM {process}') + else: + self.run_command(host, f'pkill -9 {process}') + + 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 bedc189..f7200cf 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -8,7 +8,7 @@ from multiprocessing.pool import ThreadPool from wazuh_qa_framework.generic_modules.logging.base_logger import BaseLogger -from wazuh_qa_framework.global_variables.daemons import WAZUH_ANGENT_WINDOWS_SERVICE_NAME +from wazuh_qa_framework.global_variables.daemons import WAZUH_AGENT_WINDOWS_SERVICE_NAME from wazuh_qa_framework.system.host_manager import HostManager @@ -519,7 +519,7 @@ def restart_agent(self, host): host (str): Hostname """ self.logger.debug(f'Restarting agent {host}') - service_name = WAZUH_ANGENT_WINDOWS_SERVICE_NAME if self.is_windows(host) else 'wazuh-agent' + 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, 'restarted') self.logger.debug(f'Agent {host} restarted successfully') @@ -558,7 +558,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}') @@ -569,50 +569,43 @@ 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, type='gracefully'): - """Stop agent + def stop_agent(self, host, gracefully=True): + """Stop agent. Args: - host (str): Hostname - type (str): Type of stop, it can be gracefully or ungracefully. Defaults to gracefully + host (str): Hostname. + gracefully (bool): Stop gracefully. Defaults to True. """ - self.logger.debug(f'Stopping agent {type} {host}') - service_name = WAZUH_ANGENT_WINDOWS_SERVICE_NAME if self.is_windows(host) else 'wazuh-agent' + 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): - if type == 'gracefully': - if self.is_windows(host): - self.control_service(host, service_name, 'stopped') - self.logger.debug(f'Agent {host} stopped successfully') - else: - self.control_service(host, service_name, 'stopped', become=True) - self.logger.debug(f'Agent {host} stopped successfully') - - if type == 'ungracefully': - if self.is_windows(host): - self.run_command(host, 'taskkill /F /IM wazuh-agent.exe') - self.logger.debug(f'Agent {host} stopped successfully') - else: - self.run_command(host, "pkill -9 wazuh-agent", become=True) - self.logger.debug(f'Agent {host} stopped successfully') + if gracefully: + self.control_service(host, service_name, 'stopped') + + if not gracefully: + process = 'wazuh-agent.exe' if self.is_windows(host) else 'wazuh-agent' + self.stop_ungracefully_process(host, process) + + self.logger.info(f'Agent {host} stopped {type}') + else: - raise ValueError(f'Host {host} is not an agent') + raise ValueError(f'Host {host} is not an Agent') - def stop_agents(self, agent_list=None, parallel=True, type='gracefully'): - """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. - type (str): Type of stop, it can be gracefully or ungracefully. Defaults to gracefully + gracefully (bool): Stop gracefully. Defaults to True. """ - self.logger.info(f'Stopping agents {type}: {agent_list}') if parallel: - stop_agent_with_type = partial(self.stop_agent, type=type) + stop_agent_with_type = partial(self.stop_agent, gracefully=gracefully) self.pool.map(stop_agent_with_type, agent_list) else: for agent in agent_list: - self.stop_agent(agent, type) - self.logger.info(f'Agents stopped successfully: {agent_list}') + self.stop_agent(agent, gracefully) def stop_manager(self, host): """Stop manager @@ -649,7 +642,7 @@ def start_agent(self, host): host (str): Hostname """ self.logger.debug(f'Starting agent {host}') - service_name = WAZUH_ANGENT_WINDOWS_SERVICE_NAME if self.is_windows(host) else 'wazuh-agent' + 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, 'started') self.logger.debug(f'Agent {host} started successfully') @@ -744,11 +737,11 @@ def stop_environment(self, parallel=True): self.pool.map(self.stop_agent, agent_list) else: self.logger.info(message='Stopping environment: Managers') - for manager in manager_list: + for manager in agent_list: self.stop_manager(manager) self.logger.info(message='Stopping environment: Agents') - for agent in manager_list: + for agent in agent_list: self.stop_agent(agent) self.logger.info('Stopping environment') From 1d588e97672da058370e37243e615872fc0d0799 Mon Sep 17 00:00:00 2001 From: Antonio Date: Fri, 28 Jul 2023 10:55:56 +0200 Subject: [PATCH 4/6] fix(#61): fixes done after review --- src/wazuh_qa_framework/system/wazuh_handler.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index f7200cf..f912643 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -580,17 +580,16 @@ def stop_agent(self, host, gracefully=True): 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): - if gracefully: - self.control_service(host, service_name, 'stopped') - 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') + raise ValueError(f'Host {host} is not an agent') def stop_agents(self, agent_list=None, parallel=True, gracefully=True): """Stop agents. @@ -606,6 +605,7 @@ def stop_agents(self, agent_list=None, parallel=True, gracefully=True): else: for agent in agent_list: self.stop_agent(agent, gracefully) + self.logger.info(f'Agents stopped successfully: {agent_list}') def stop_manager(self, host): """Stop manager @@ -737,7 +737,7 @@ def stop_environment(self, parallel=True): self.pool.map(self.stop_agent, agent_list) else: self.logger.info(message='Stopping environment: Managers') - for manager in agent_list: + for manager in manager_list: self.stop_manager(manager) self.logger.info(message='Stopping environment: Agents') From 1ebf89058de3b4a0eae192c0ede83285e0480140 Mon Sep 17 00:00:00 2001 From: Antonio Date: Thu, 16 Nov 2023 18:48:46 +0100 Subject: [PATCH 5/6] fix(#60): Fixes after review --- src/wazuh_qa_framework/system/host_manager.py | 4 ++-- src/wazuh_qa_framework/system/wazuh_handler.py | 5 ++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index c5916d5..7124830 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -643,9 +643,9 @@ def stop_ungracefully_process(self, host, process): """ try: if self.is_windows(host): - self.run_command(host, f'taskkill /F /IM {process}') + self.run_command(host, f"taskkill /F /IM {process}") else: - self.run_command(host, f'pkill -9 {process}') + self.run_command(host, f"pkill -9 {process}") except Exception as e: raise Exception(f"An error occurred while stopping the process. {e}") diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index d46bd37..0c2183d 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -707,7 +707,7 @@ def stop_agent(self, host, gracefully=True): else: self.control_service(host, service_name, 'stopped') - self.logger.info(f'Agent {host} stopped {type}') + self.logger.info(f"Agent {host} stopped {type}") else: raise ValueError(f"Host {host} is not an agent") @@ -721,8 +721,7 @@ def stop_agents(self, agent_list=None, parallel=True, gracefully=True): gracefully (bool): Stop gracefully. Defaults to True. """ if parallel: - stop_agent_with_type = partial(self.stop_agent, gracefully=gracefully) - self.pool.map(stop_agent_with_type, agent_list) + self.pool.starmap(self.stop_agent, [(agent, gracefully) for agent in agent_list]) else: for agent in agent_list: self.stop_agent(agent, gracefully) From 96dda29fbd72775eddafe2d1a9203961d14b201f Mon Sep 17 00:00:00 2001 From: Antonio Date: Wed, 22 Nov 2023 11:53:18 +0100 Subject: [PATCH 6/6] fix(#61): remove library and add become --- src/wazuh_qa_framework/system/host_manager.py | 8 +++++--- src/wazuh_qa_framework/system/wazuh_handler.py | 1 - 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index 7124830..c6aa731 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -634,18 +634,20 @@ def install_package(self, host, package_name, become=None, ignore_errors=False): return result - def stop_ungracefully_process(self, host, process): + 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}") + self.run_command(host, f"taskkill /F /IM {process}", become=become) else: - self.run_command(host, f"pkill -9 {process}") + 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}") diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index 0c2183d..51a6b6e 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -5,7 +5,6 @@ import os import re import yaml -from functools import partial from multiprocessing.pool import ThreadPool from wazuh_qa_framework.generic_modules.logging.base_logger import BaseLogger