Skip to content

Commit

Permalink
fix(#46): Adapting methods after merge
Browse files Browse the repository at this point in the history
  • Loading branch information
@pro-akim
pro-akim committed Jun 15, 2023
1 parent 007d5ac commit 75bb26a
Showing 1 changed file with 26 additions and 53 deletions.
79 changes: 26 additions & 53 deletions src/wazuh_qa_framework/system/wazuh_handler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
"""
Expand All @@ -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
Expand Down Expand Up @@ -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')
Expand Down Expand Up @@ -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')
Expand Down Expand Up @@ -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')
Expand All @@ -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')
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -865,13 +839,15 @@ 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)
else:
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')
Expand All @@ -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
Expand Down

0 comments on commit 75bb26a

Please sign in to comment.