Skip to content

Commit

Permalink
refactor(#46): New changes after review
Browse files Browse the repository at this point in the history
  • Loading branch information
@pro-akim
pro-akim committed Jun 15, 2023
1 parent 172235f commit 508698d
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 123 deletions.
61 changes: 0 additions & 61 deletions src/wazuh_qa_framework/system/host_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -641,64 +641,3 @@ def append_block_in_file(self, host, path, block, become=None, ignore_errors=Fal
raise Exception(f"Error inserting a block in file {path} on host {host}: {result}")

return result

def get_api_token(self, host, user='wazuh', password='wazuh', auth_context=None, port=55000, check=False):
"""Return an API token for the specified user.
Args:
host (str): Hostname.
user (str, optional): API username. Default `wazuh`
password (str, optional): API password. Default `wazuh`
auth_context (dict, optional): Authorization context body. Default `None`
port (int, optional): API port. Default `55000`
check (bool, optional): Ansible check mode("Dry Run"),
by default it is enabled so no changes will be applied. Default `False`
Returns:
API token (str): Usable API token.
"""
login_endpoint = '/security/user/authenticate'
login_method = 'POST'
login_body = ''
if auth_context is not None:
login_endpoint = '/security/user/authenticate/run_as'
login_body = 'body="{}"'.format(json.dumps(auth_context).replace('"', '\\"').replace(' ', ''))

try:
token_response = self.get_host(host).ansible('uri', f"url=https://localhost:{port}{login_endpoint} "
f"user={user} password={password} "
f"method={login_method} {login_body} validate_certs=no "
f"force_basic_auth=yes",
check=check)
return token_response['json']['data']['token']
except KeyError:
raise KeyError(f'Failed to get token: {token_response}')

def make_api_call(self, host, port=55000, method='GET', endpoint='/', request_body=None, token=None, check=False):
"""Make an API call to the specified host.
Args:
host (str): Hostname.
port (int, optional): API port. Default `55000`
method (str, optional): Request method. Default `GET`
endpoint (str, optional): Request endpoint. It must start with '/'.. Default `/`
request_body ( dict, optional) : Request body. Default `None`
token (str, optional): Request token. Default `None`
check ( bool, optional): Ansible check mode("Dry Run"), by default it is enabled so no changes will be
applied. Default `False`
Returns:
API response (dict) : Return the response in JSON format.
"""
request_body = 'body="{}"'.format(
json.dumps(request_body).replace('"', '\\"').replace(' ', '')) if request_body else ''

token = self.get_api_token(host, user='wazuh', password='wazuh', auth_context=None, port=55000, check=False)

headers = {'Authorization': f'Bearer {token}'}
if request_body:
headers['Content-Type'] = 'application/json'

return self.get_host(host).ansible('uri', f'url="https://localhost:{port}{endpoint}" '
f'method={method} headers="{headers}" {request_body} '
f'validate_certs=no', check=check)
119 changes: 57 additions & 62 deletions src/wazuh_qa_framework/system/wazuh_handler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
from multiprocessing.pool import ThreadPool

from wazuh_qa_framework.system.host_manager import HostManager
from wazuh_qa_framework.wazuh_components.api.wazuh_api import WazuhAPI
from wazuh_qa_framework.wazuh_components.api.wazuh_api_request import WazuhAPIRequest

DEFAULT_INSTALL_PATH = {
'linux': '/var/ossec',
Expand Down Expand Up @@ -52,7 +54,7 @@ def get_archives_directory_path(custom_installation_path=None):

def get_logs_directory_path(custom_installation_path=None, os_host='linux'):
installation_path = custom_installation_path if custom_installation_path else DEFAULT_INSTALL_PATH[os_host]
return installation_path + '\\logs' if os_host == 'windows' else os.path.join(installation_path, 'logs')
return installation_path if os_host == 'windows' else os.path.join(installation_path, 'logs')


def get_shared_directory_path(custom_installation_path=None, os_host='linux'):
Expand Down Expand Up @@ -123,7 +125,7 @@ def get_wazuh_file_path(custom_installation_path=None, os_host='linux', file_nam
'custom_rule_directory': {
'files': ['local_rules.xml'],
'path_calculator': lambda filename: os.path.join(get_custom_rules_directory_path(installation_path),
filename)
filename)
},
'group_configuration': {
'files': ['agent.conf'],
Expand Down Expand Up @@ -495,32 +497,22 @@ def get_agents_info(self):
"""
pass

def get_agents_id(self, agent_list):
"""Get agent ids

def get_agent_id(self, host, agent):
"""Get agent id
Args:
agents_list (_type_, agents_list): Agents list.
host (_type_, str): Ansible host name.
agent (_type_, str): Agent name.
Return:
dict: agent_ids
str: agent_id
"""
# Getting hostnames
host_names = []
for agent in agent_list:
host_names.append(self.run_command(agent, 'hostname')[1])
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')

# Getting id - hostnames from manager
agent_control = self.run_command('manager1', '/var/ossec/bin/manage_agents -l', True)[1]
return None

# Creating id_list from hostnames
agent_ids = []
for hostname in host_names:
hostname = hostname.replace('\r', '').replace('\n', '')
for line in agent_control.split('\n'):
if 'Name: ' + hostname in line:
id_value = line.split(',')[0].split(': ')[1].strip()
agent_ids.append(id_value)
break

return agent_ids

def restart_manager(self, host):
"""Restart manager
Expand Down Expand Up @@ -580,21 +572,21 @@ def clean_client_keys(self, hosts=None):
"""
pass

def clean_logs(self, hosts):
def clean_logs(self, host):
"""Remove host logs
Args:
hosts (_type_, hosts): host list.
host (_type_, str): Host.
"""
# Clean ossec.log and and cluster.log
for host in hosts:
logs_path = self.get_logs_directory_path(host)
if self.get_host_variables(host)['os_name'] == 'windows':
self.truncate_file(host, f'{logs_path}/ossec.log', recreate=True, become=False, ignore_errors=False)
else:
self.truncate_file(host, f'{logs_path}/ossec.log', recreate=True, become=True, ignore_errors=False)
host_type = self.get_host_variables(host).get('type')
if 'master' == host_type or 'worker' == host_type:
self.truncate_file(host, f'{logs_path}/cluster.log', recreate=True, become=True, ignore_errors=False)
# Clean ossec.log, api.log and cluster.log
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)
else:
self.truncate_file(host, f'{logs_path}/ossec.log', recreate=True, become=True, ignore_errors=False)
host_type = self.get_host_variables(host).get('type')
if 'master' == host_type or 'worker' == host_type:
self.truncate_file(host, f'{logs_path}/api.log', recreate=True, become=True, ignore_errors=False)
self.truncate_file(host, f'{logs_path}/cluster.log', recreate=True, become=True, ignore_errors=False)

def clean_agents(self, agents=None):
"""Stop agents, remove them from manager and clean their client keys
Expand All @@ -603,20 +595,19 @@ def clean_agents(self, agents=None):
"""
pass

def restart_agents(self, agent_list):
def restart_agent(self, agent):
"""Restart agents
Args:
agents_list (_type_, agents_list): Agents list.
agent (_type_, agent): Agent.
"""
# Clean ossec.log and and cluster.log
for agent in agent_list:
if self.get_host_variables(agent).get('os_name') == 'windows':
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)
# 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 @@ -625,39 +616,43 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par
manager (str, optional): Name of manager. Defaults to None.
method (str): Method to be used to remove agents, Defaults to cmd.
parallel (str): In case that cmd method is used, it defines the use of threads for remove. Defaults to True.
logs (str): Remove logs from agents. Defaults to False.
logs (str): Remove logs (ossec.log, api.log) from agents. Defaults to False.
restart (str): Restart agents. Defaults to False.
"""
manager = 'manager1' if manager is None else manager
parallel = False if method == 'api' else parallel
if manager is None: manager = 'manager1'

# Getting agent_ids list
agent_ids = self.get_agents_id(agent_list)

# Remove agent by cmd core function
def remove_agent_cmd(id):
self.run_command(manager, f"/var/ossec/bin/manage_agents -r {id}", True)
agent_ids = []
for agent in agent_list:
agent_ids.append(self.get_agent_id(manager, agent))

# Remove processes
if method == 'cmd':
if parallel:
self.pool.map(remove_agent_cmd, agent_ids)
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:
remove_agent_cmd(id)
self.run_command(manager, f"/var/ossec/bin/manage_agents -r {id}", True)
else:
agent_string = ','.join(agent_ids)
self.make_api_call('manager1', port=55000, method='DELETE',
endpoint=f'/agents?pretty=true&older_than=0s&agents_list={agent_string}&status=all',
request_body=None, token=None, check=False)
endpoint=f'/agents?pretty=true&older_than=0s&agents_list={agent_string}&status=all'
request = WazuhAPIRequest(endpoint=endpoint, method='DELETE')
request.send(WazuhAPI(address = self.get_host_variables(manager)['ip']))

# Remove logs
if logs:
self.clean_logs(agent_list)
if logs and parallel == False:
for agent in agent_list:
self.clean_logs(agent)
if logs and parallel == True:
self.pool.map(self.clean_logs, agent_list)

# Restarting agents
if restart:
self.restart_agents(agent_list)
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)

def stop_manager(self, manager):
"""Stop manager
Expand Down

0 comments on commit 508698d

Please sign in to comment.