Merge pull request #72 from wazuh/61-stop-agents-ungracefully

Stopping agents ungracefully feature

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.
 

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')