From 75bb26a35b44e7722d3b15d854d33f22aeacb758 Mon Sep 17 00:00:00 2001 From: Antonio Date: Thu, 15 Jun 2023 16:08:07 +0200 Subject: [PATCH] fix(#46): Adapting methods after merge --- .../system/wazuh_handler.py | 79 ++++++------------- 1 file changed, 26 insertions(+), 53 deletions(-) diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index 278ecdb..a406552 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -505,15 +505,20 @@ def get_agents_info(self): """ pass -<<<<<<< HEAD -======= - def get_agents_id(self, agents_list=None): - """Get agents id - - Returns: - List: Agents id list + def get_agent_id(self, host, agent): + """Get agent id + Args: + host (_type_, str): Ansible host name. + agent (_type_, str): Agent name. + Return: + str: agent_id """ - pass + host_list = WazuhAPI(address = self.get_host_variables(host)['ip']).list_agents()['affected_items'] + for host in host_list: + if host.get('ip') == self.get_host_variables(agent)['ip']: + return host.get('id') + + return None def restart_agent(self, host): """Restart agent @@ -530,26 +535,7 @@ def restart_agent(self, host): raise ValueError(f'Host {host} is not an agent') def restart_agents(self, agent_list=None, parallel=True): - """Restart list of agents ->>>>>>> system-refactor - - def get_agent_id(self, host, agent): - """Get agent id - Args: -<<<<<<< HEAD - host (_type_, str): Ansible host name. - agent (_type_, str): Agent name. - Return: - str: agent_id - """ - host_list = WazuhAPI(address = self.get_host_variables(host)['ip']).list_agents()['affected_items'] - for host in host_list: - if host.get('ip') == self.get_host_variables(agent)['ip']: - return host.get('id') - - return None - -======= + """ Restart list of agents agent_list (list, optional): Agent list. Defaults to None. parallel (bool, optional): Parallel execution. Defaults to True. """ @@ -560,7 +546,6 @@ def get_agent_id(self, host, agent): for agent in agent_list: self.restart_agent(agent) self.logger.info(f'Agents restarted successfully: {agent_list}') ->>>>>>> system-refactor def restart_manager(self, host): """Restart manager @@ -597,7 +582,7 @@ def stop_agent(self, host): host (str): Hostname """ self.logger.debug(f'Stopping 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, 'stopped') self.logger.debug(f'Agent {host} stopped successfully') @@ -654,7 +639,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') @@ -749,11 +734,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 agent_list: self.stop_agent(agent) self.logger.info('Stopping environment') @@ -776,11 +761,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') @@ -815,6 +800,7 @@ def clean_logs(self, host): host (_type_, str): Host. """ # Clean ossec.log, api.log and cluster.log + self.logger.info(f'Removing {host} logs') logs_path = self.get_logs_directory_path(host) if self.is_windows(host): self.truncate_file(host, f'{logs_path}/ossec.log', recreate=True, become=False, ignore_errors=False) @@ -832,19 +818,7 @@ def clean_agents(self, agents=None): """ pass - def restart_agent(self, agent): - """Restart agents - Args: - agent (_type_, agent): Agent. - """ - # Restart agent - if self.is_windows(agent): - self.run_command(agent, f"NET STOP WazuhSvc", become=False, ignore_errors=False) - self.run_command(agent, f"NET START WazuhSvc", become=False, ignore_errors=False) - else: - self.run_command(agent, f"service wazuh-agent restart", become=True, ignore_errors=False) - - def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', parallel=True , logs=False, + def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', parallel=True, logs=False, restart=False): """Remove agents from manager @@ -865,6 +839,7 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par # Remove processes if method == 'cmd': + self.logger.info(f'Removing agents {agent_list} using cmd') if parallel: self.pool.map(lambda id: self.run_command(manager, f"/var/ossec/bin/manage_agents -r {id}", True), agent_ids) @@ -872,6 +847,7 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par for id in agent_ids: self.run_command(manager, f"/var/ossec/bin/manage_agents -r {id}", True) else: + self.logger.info(f'Removing agents {agent_list} using API') agent_string = ','.join(agent_ids) endpoint=f'/agents?pretty=true&older_than=0s&agents_list={agent_string}&status=all' request = WazuhAPIRequest(endpoint=endpoint, method='DELETE') @@ -885,11 +861,8 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par self.pool.map(self.clean_logs, agent_list) # Restarting agents - if restart and parallel == False: - for agent in agent_list: - self.restart_agent(agent) - if restart and parallel == True: - self.pool.map(self.restart_agent, agent_list) + if restart: + self.restart_agents(agent_list, parallel=parallel) def get_managers(self): """Get environment managers names