Skip to content

Commit

Permalink
fix(#60): Small fixes after linting, get_agent_id parameters
Browse files Browse the repository at this point in the history
  • Loading branch information
@pro-akim
pro-akim committed Aug 23, 2023
1 parent aea23a1 commit 33af533
Showing 1 changed file with 14 additions and 14 deletions.
28 changes: 14 additions & 14 deletions src/wazuh_qa_framework/system/wazuh_handler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -509,7 +509,7 @@ def get_agents_info(self):
"""
pass

def get_agent_id(self, agent_name):
def get_agent_id(self, manager, agent_name):
"""Get agent id
Returns:
Expand All @@ -518,7 +518,7 @@ def get_agent_id(self, agent_name):
agent_ip = self.get_host_variables(agent_name).get('ip')
endpoint = f'/agents'
request = WazuhAPIRequest(endpoint=endpoint, method='GET')
for item in request.send(WazuhAPI(address=self.get_host_variables('manager1')['ip'])).data['affected_items']:
for item in request.send(WazuhAPI(address=self.get_host_variables(manager)['ip'])).data['affected_items']:
if item.get('ip') == agent_ip:
return item.get('id')

Expand All @@ -533,7 +533,7 @@ def get_agent_name_from_ip(self, agent_ip):
if self.get_host_variables(host).get('ip') == agent_ip:
return self.get_host_variables(host).get('inventory_hostname_short')

def get_agents_id(self, agents_list=None):
def get_agents_id(self, manager, agents_list=None):
"""Get agents id
Returns:
Expand All @@ -544,7 +544,7 @@ def get_agents_id(self, agents_list=None):
agent_ip = self.get_host_variables(agent).get('ip')
endpoint = f'/agents'
request = WazuhAPIRequest(endpoint=endpoint, method='GET')
items = request.send(WazuhAPI(address=self.get_host_variables('manager1')['ip'])).data['affected_items']
items = request.send(WazuhAPI(address=self.get_host_variables(manager)['ip'])).data['affected_items']
for item in items:
if item.get('ip') == agent_ip:
result_id.append(item.get('id'))
Expand Down Expand Up @@ -955,10 +955,10 @@ def create_group(self, manager, group_name, method='api', check_group=True):
else:
self.logger.info(f'Creating group {group_name} from {manager} using {method.upper()}')
if method == 'cmd':
self.run_command(manager, f"{get_bin_directory_path()}/agent_groups -q -a -g {group_name}")
self.run_command(manager, f'{get_bin_directory_path()}/agent_groups -q -a -g {group_name}')
elif method == 'api':
endpoint = '/groups'
request = WazuhAPIRequest(endpoint=endpoint, payload={"group_id": group_name}, method='POST')
request = WazuhAPIRequest(endpoint=endpoint, payload={'group_id': group_name}, method='POST')
request.send(WazuhAPI(address=self.get_host_variables(manager)['ip']))

def delete_group(self, manager, group_name, method='cmd', check_group=True):
Expand All @@ -974,15 +974,15 @@ def delete_group(self, manager, group_name, method='cmd', check_group=True):
else:
self.logger.info(f'Removing group {group_name} from {manager} using {method.upper()} method')
if method == 'cmd':
self.run_command(manager, f"{get_bin_directory_path()}/agent_groups -q -r -g {group_name}")
self.run_command(manager, f'{get_bin_directory_path()}/agent_groups -q -r -g {group_name}')

if method == 'api':
endpoint = f'/groups?groups_list={group_name}'
request = WazuhAPIRequest(endpoint=endpoint, method='DELETE')
request.send(WazuhAPI(address=self.get_host_variables(manager)['ip']))

if method == 'folder':
self.run_command(manager, f"rm -rf {get_shared_directory_path()}/{group_name}")
self.run_command(manager, f'rm -rf {get_shared_directory_path()}/{group_name}')

def assign_agent_group(self, manager, agent_name, group_name, method='api', check_group=True):
"""Assign an agent to a group.
Expand All @@ -998,11 +998,11 @@ def assign_agent_group(self, manager, agent_name, group_name, method='api', chec
else:
self.logger.info(f'Assign agent {agent_name} to group {group_name} from {manager} using {method.upper()}')
if method == 'cmd':
self.run_command(manager, (f"{get_bin_directory_path()}/agent_groups -q -a "
f"-i {self.get_agent_id(agent_name)} -g {group_name}"))
self.run_command(manager, (f'{get_bin_directory_path()}/agent_groups -q -a '
f'-i {self.get_agent_id(manager, agent_name)} -g {group_name}'))

if method == 'api':
endpoint = f'/agents/{self.get_agent_id(agent_name)}/group/{group_name}'
endpoint = f'/agents/{self.get_agent_id(manager, agent_name)}/group/{group_name}'
request = WazuhAPIRequest(endpoint=endpoint, method='PUT')
request.send(WazuhAPI(address=self.get_host_variables(manager)['ip']))

Expand All @@ -1018,7 +1018,7 @@ def assign_agents_group(self, manager, list_agent_names, group_name, method='api
"""
if parallel:
self.pool.map(lambda agent: self.assign_agent_group(manager, agent, group_name, method=method,
check_group=check_group ),list_agent_names)
check_group=check_group ), list_agent_names)
else:
for agent in list_agent_names:
self.logger.info(f'Assigning agent {agent} from group {group_name} from {manager}')
Expand All @@ -1037,7 +1037,7 @@ def unassign_agent_group(self, manager, agent_name, group_name, check_group=True
if check_group and not self.check_agent_group(manager, agent_name, group_name):
self.logger.info(f'{agent_name} is not assigned to {group_name}')
else:
endpoint = f'/agents/{self.get_agent_id(agent_name)}/group/{group_name}'
endpoint = f'/agents/{self.get_agent_id(manager, agent_name)}/group/{group_name}'
request = WazuhAPIRequest(endpoint=endpoint, method='DELETE')
request.send(WazuhAPI(address=self.get_host_variables(manager)['ip']))

Expand All @@ -1054,5 +1054,5 @@ def unassign_agents_group(self, manager, list_agent_names, group_name, check_gro
list_agent_names)

else:
for agent in self.get_agents_id(list_agent_names):
for agent in self.get_agents_id(manager, list_agent_names):
self.unassign_agent_group(manager, agent, group_name, check_previous=check_group)

0 comments on commit 33af533

Please sign in to comment.