From 508698d1b664d38a5b033d782ae0a839951eb646 Mon Sep 17 00:00:00 2001 From: Antonio Date: Thu, 15 Jun 2023 11:18:53 +0200 Subject: [PATCH] refactor(#46): New changes after review --- src/wazuh_qa_framework/system/host_manager.py | 61 --------- .../system/wazuh_handler.py | 119 +++++++++--------- 2 files changed, 57 insertions(+), 123 deletions(-) diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index ee30a64..ec26352 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -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) diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index be59a2c..da6722b 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -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', @@ -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'): @@ -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'], @@ -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 @@ -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 @@ -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 @@ -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