From 00d95dc9dd4fff12c2aab7f12001501d724576f0 Mon Sep 17 00:00:00 2001 From: Antonio Date: Wed, 7 Jun 2023 13:37:12 +0200 Subject: [PATCH 1/9] feat(#46): Remove agents in construction --- src/wazuh_qa_framework/system/host_manager.py | 704 ++++++++++++++++++ .../system/wazuh_handler.py | 127 +++- 2 files changed, 799 insertions(+), 32 deletions(-) diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index cab81b5..d324504 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -6,6 +6,7 @@ import testinfra import base64 import os +import json from ansible.inventory.manager import InventoryManager from ansible.parsing.dataloader import DataLoader from ansible.vars.manager import VariableManager @@ -640,3 +641,706 @@ 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 +# Copyright (C) 2015-2021, Wazuh Inc. +# Created by Wazuh, Inc. . +# This program is free software; you can redistribute it and/or modify it under the terms of GPLv2 + +import tempfile +import testinfra +import base64 +import os +from ansible.inventory.manager import InventoryManager +from ansible.parsing.dataloader import DataLoader +from ansible.vars.manager import VariableManager + + +class HostManager: + """Remote host management interface. + + It allows to manage remote hosts using ansible inventory and testinfra framework. + + Args: + inventory_path (str): Ansible inventory path + + Attributes: + inventory_path (str): Ansible inventory path + inventory_manager (ansible.inventory.manager.InventoryManager): Ansible inventory manager + variable_manager (ansible.vars.manager.VariableManager): Ansible variable manager + """ + + def __init__(self, inventory_path): + self.inventory_path = inventory_path + + data_loader = DataLoader() + self.inventory_manager = InventoryManager(loader=data_loader, sources=inventory_path) + self.hosts_variables = {} + + variable_manager = VariableManager(loader=data_loader, inventory=self.inventory_manager) + + for host in self.inventory_manager.get_hosts(): + self.hosts_variables[host] = variable_manager.get_vars(host=self.inventory_manager.get_host(str(host))) + + def get_host(self, host): + """Get the testinfra host. + + Args: + host (str): Hostname + + Returns: + testinfra.modules.base.Ansible: Host instance from hostspec + """ + return testinfra.get_host(f"ansible://{host}?ansible_inventory={self.inventory_path}") + + def get_groups(self): + """Get the groups of the inventory. + + Returns: + list: Groups of the inventory + """ + return list(self.inventory_manager.groups.keys()) + + def get_group_hosts(self, pattern=None): + """Get all hosts from inventory that belong to a group. + + Args: + group (str): Group name + + Returns: + list: List of hosts + """ + if pattern: + return [str(host) for host in self.inventory_manager.get_hosts(pattern=pattern)] + else: + return [str(host) for host in self.inventory_manager.get_hosts()] + + def get_host_variables(self, host): + """Get the variables of the specified host. + + Args: + host (str): Hostname + + Returns: + testinfra.modules.base.Ansible: Host instance from hostspec + """ + inventory_manager_host = self.inventory_manager.get_host(host) + + return self.hosts_variables[inventory_manager_host] + + def collect_host_ansible_facts(self, host): + """Get the ansible facts of the specified host. + + Args: + host (str): Hostname + + Returns: + str: OS of the host + """ + testinfra_host = self.get_host(host) + + return testinfra_host.ansible("setup") + + def collect_host_os(self, host): + """Get the OS of the specified host. + + Args: + host: Hostname + + Returns: + tuple: Hostname, Major version, Distribution version. Example: ('CentOS', '7', '7.6.1810') + """ + ansible_facts = self.collect_host_ansible_facts(host) + + return (ansible_facts['ansible_facts']['ansible_distribution'], + ansible_facts['ansible_facts']['ansible_distribution_major_version'], + ansible_facts['ansible_facts']['ansible_distribution_version']) + + def collect_host_ips(self, host): + """Get the host IPs + + Args: + host (str): Hostname + + Returns: + dict: IPs of the host (ipv4 and ipv6). Example: {'ipv4': ['172.31.5.209'], 'ipv6': ['fe80::f::fef4:bb6d']} + """ + ansible_facts = self.collect_host_ansible_facts(host) + + return {'ipv4': ansible_facts['ansible_facts']['ansible_all_ipv4_addresses'], + 'ipv6': ansible_facts['ansible_facts']['ansible_all_ipv6_addresses']} + + def collect_host_interfaces(self, host): + """Get the interfaces of the specified host. + + Args: + host (str): Hostname + + Returns: + list: Interfaces of the host. Example ['lo', 'eth0'] + """ + ansible_facts = self.collect_host_ansible_facts(host) + + return ansible_facts['ansible_facts']['ansible_interfaces'] + + def check_connection(self, host): + """Check if the host is reachable. + + Args: + host (str): Hostname + + Returns: + bool: True if the host is reachable, False otherwise + """ + testinfra_host = self.get_host(host) + ansible_command = 'win_ping' if self.get_host_variables(host)['os_name'] == 'windows' else 'ping' + return testinfra_host.ansible(ansible_command, check=False)['ping'] == 'pong' + + def copy_file(self, host, src_path, dest_path, remote_src=False, become=None, ignore_errors=False): + """Move from src_path to the desired location dest_path for the specified host. + + Args: + host (str): Hostname + src_path (str): Source path + dest_path (str): Destination path + remote_src (bool): If True, the file is assumed to live on the remote machine, not the controller. + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool): Ignore errors + + Returns: + dict: Result of the command execution + + Raises: + Exception: If the command execution fails + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + ansible_command = 'win_copy' if self.get_host_variables(host)['os_name'] == 'windows' else 'copy' + remote_source = 'yes' if remote_src else 'no' + + command_parameters = f"src='{src_path}' dest='{dest_path}' remote_src={remote_source}" + result = testinfra_host.ansible(ansible_command, command_parameters, check=False, become=become) + + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error moving file from {src_path} to {dest_path} on host {host}: {result}") + + return result + + def get_file_content(self, host, path, become=None, ignore_errors=False): + """Read a file from the specified host. + + Args: + host (str): Hostname + path (str): File path + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool): Ignore errors + + Returns: + str: File content + + Raises: + Exception: If the file cannot be read + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + result = testinfra_host.ansible("slurp", f"src='{path}'", check=False, become=become) + + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error reading file {path} on host {host}: {result}") + + return base64.b64decode(result['content']).decode('utf-8') + + def synchronize_linux_directory(self, host, dest_path, src_path=None, filesystem=None, become=None, + ignore_errors=False): + """Create a file structure on the specified host. + Not supported on Windows. + + Args: + host (str): Hostname + dest_path (str): Destination path + filesystem (dict): File structure + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Result of the command execution + + Raises: + Exception: If the command execution fails + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + + ansible_command = 'synchronize' + + if filesystem: + tmp_directory = tempfile.TemporaryDirectory() + directory_path = os.path.join(tmp_directory.name, filesystem['directory_name']) + os.mkdir(directory_path) + src_path = directory_path + + for file in filesystem['files']: + file_path = f"{directory_path}/{file['filename']}" + with open(file_path, 'w') as file_operator: + file_operator.write(file['content']) + + result = testinfra_host.ansible(ansible_command, f"src='{src_path}' dest='{dest_path}'", check=False, + become=become) + + if (result['rc'] != 0 or not result) and not ignore_errors: + raise Exception(f"Error creating file structure on host {host}: {result}") + + return result + + def truncate_file(self, host, file_path, recreate=True, become=None, ignore_errors=False): + """Truncate a file from the specified host. + + Args: + host (str): Hostname + file_path (str): File path + recreate (bool, optional): Recreate file. Defaults to True. + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the file cannot be truncated + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + result = None + + if recreate: + ansible_command = 'win_copy' if self.get_host_variables(host)['os_name'] == 'windows' else 'copy' + result = testinfra_host.ansible(ansible_command, f"dest='{file_path}' content=''", check=False, + become=become) + else: + ansible_command = 'win_file' if self.get_host_variables(host)['os_name'] == 'windows' else 'file' + result = testinfra_host.ansible(ansible_command, f"path='{file_path}' state=touch", check=False, + become=become) + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error truncating file {file_path} on host {host}: {result}") + + return result + + def remove_file(self, host, file_path, become=None, ignore_errors=False): + """Remove a file from the specified host. + + Args: + host (str): Hostname + file_path (str): File path + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the file cannot be removed + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + ansible_command = 'win_file' if self.get_host_variables(host)['os_name'] == 'windows' else 'file' + result = testinfra_host.ansible(ansible_command, f"path='{file_path}' state=absent", check=False, become=become) + + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error removing file {file_path} on host {host}: {result}") + + return result + + def modify_file_content(self, host, path, content, become=None, ignore_errors=False): + """Create a file with a specified content and copies it to a path. + + Args: + host (str): Hostname + path (str): path for the file to create and modify + content (str, bytes): content to write into the file + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the file cannot be modified + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + tmp_file = tempfile.NamedTemporaryFile() + with open(tmp_file.name, 'w+') as tmp: + tmp.write(content) + + result = self.copy_file(host, src_path=tmp_file.name, dest_path=path, become=become) + + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error modifying file {path} on host {host}: {result}") + + return result + + def create_file(self, host, path, content, directory=False, owner=None, group=None, mode=None, become=None, + ignore_errors=False): + """Create a file with a specified content and copies it to a path. + + Args: + host (str): Hostname + path (str): path for the file to create and modify + content (str, bytes): content to write into the file + owner (str): owner of the file + group (str): group of the file + mode (str): mode of the file + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the file cannot be created + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + tmp_file = tempfile.NamedTemporaryFile() + with open(tmp_file.name, 'w+') as tmp: + tmp.write(content) + + ansible_command = 'win_copy' if self.get_host_variables(host)['os_name'] == 'windows' else 'copy' + + ansible_parameters = f"src='{tmp_file.name}' dest='{path}'" + ansible_parameters += f" owner={owner}" if owner else '' + ansible_parameters += f" group={group}" if group else '' + ansible_parameters += f" mode={mode}" if mode else '' + ansible_parameters += f' state=directory' if directory else '' + + result = testinfra_host.ansible(ansible_command, ansible_parameters, check=False, become=become) + + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error creating file {path} on host {host}: {result}") + return result + + def control_service(self, host, service, state, become=None, ignore_errors=False): + """Control a service on a host. + + Args: + host (str): Hostname + service (str): Service name + state (str): Service state + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the service cannot be controlled + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + ansible_command = 'win_service' if self.get_host_variables(host)['os_name'] == 'windows' else 'service' + + result = testinfra_host.ansible(ansible_command, f"name={service} state={state}", check=False, become=become) + + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error controlling service {service} on host {host}: {result}") + + return result + + def run_command(self, host, cmd, become=None, ignore_errors=False): + """Run a command on a host. + + Args: + host (str): Hostname + cmd (str): Command to run + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the command cannot be run + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + ansible_command = 'win_command' if self.get_host_variables(host)['os_name'] == 'windows' else 'command' + + result = testinfra_host.ansible(ansible_command, f"{cmd}", check=False, become=become) + rc, stdout = result.get('rc', 1), result.get('stdout', '') + + if rc != 0 and not ignore_errors: + raise Exception(f"Error running command '{cmd}' on host {host}: {result}") + + return rc, stdout + + def run_shell(self, host, cmd, become=None, ignore_errors=False): + """Run a shell command on a host. + The difference with run_command is that here, shell symbols like &, |, etc. are interpreted. + + Args: + host (str): Hostname + cmd (str): Command to run + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the command cannot be run + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + rc = None + stdout = None + + ansible_command = 'win_shell' if self.get_host_variables(host)['os_name'] == 'windows' else 'shell' + + result = testinfra_host.ansible(ansible_command, f"{cmd}", check=False, become=become) + + rc, stdout = result['rc'], result['stdout'] + + if rc != 0 and not ignore_errors: + raise Exception(f"Error running command {cmd} on host {host}: {result}") + + return rc, stdout + + def find_files(self, host, path, pattern, recurse=False, use_regex=False, become=None, ignore_errors=False): + """Search and return information of a file inside a path. + + Args: + host (str): Hostname + path (str): Path in which to search for the file that matches the pattern. + pattern (str): Restrict the files to be returned to those whose basenames match the pattern specified. + recurse (bool): If target is a directory, recursively descend into the directory looking for files. + use_regex (bool): If no, the patterns are file globs (shell), if yes, they are python regexes. + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + Files (list): List of found files. + + Raises: + Exception: If the command cannot be run + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + test_infra_host = self.get_host(host) + ansible_command = 'win_find' if self.get_host_variables(host)['os_name'] == 'windows' else 'find' + ansible_pattern_arguments = 'patterns' if self.get_host_variables(host)['os_name'] == 'windows' else 'pattern' + + result = test_infra_host.ansible(ansible_command, f"paths={path} {ansible_pattern_arguments}='{pattern}' \ + recurse={recurse} use_regex={use_regex}", + become=become, check=False) + + if 'files' not in result and not ignore_errors: + raise Exception(f"Error finding file {path} on host {host}: {result}") + + return result['files'] + + def get_file_stats(self, host, path, become=None, ignore_errors=False): + """Retrieve file or file system status. + + Args: + host (str): Hostname. + path (str): The full path of the file/object to get the facts of. + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result. + + Raises: + Exception: If the command cannot be run. + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + + if self.get_host_variables(host)['os_name'] == 'windows': + ansible_command = 'ansible.windows.win_stat' + else: + ansible_command = 'stat' + + result = testinfra_host.ansible(ansible_command, f"path='{path}'", check=False, become=become) + + if 'stat' not in result and not ignore_errors: + raise Exception(f"Error getting stats of {path} on host {host}: {result}") + + return result + + def install_package(self, host, package_name, become=None, ignore_errors=False): + """Install a package on a host. + + Args: + host (str): Hostname + package_name (str): Package to install + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the install cannot be run + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + ansible_arguments = f"name={package_name} state=present" + ansible_command = 'win_chocolatey' if self.get_host_variables(host)['os_name'] == 'windows' else 'package' + + result = testinfra_host.ansible(ansible_command, ansible_arguments, + check=False, become=become) + + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error installing package {package_name} on host {host}: {result}") + + return result + + def uninstall_package(self, host, package_name, become=None, ignore_errors=False): + """Uninstall a package on a host. + + Args: + host (str): Hostname + package_name (str): Package to uninstall + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the uninstall cannot be run + """ + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + ansible_arguments = f"name={package_name} state=absent" + ansible_command = 'win_chocolatey' if self.get_host_variables(host)['os_name'] == 'windows' else 'package' + + result = testinfra_host.ansible(ansible_command, ansible_arguments, + check=False, become=become) + + if result.get('msg', None) and not ignore_errors: + raise Exception(f"Error installing package {package_name} on host {host}: {result}") + + return result + + def append_block_in_file(self, host, path, block, become=None, ignore_errors=False): + """Append a text block in file + + Args: + host (str): Hostname + path (str): path for the file to insert a block + block (str, bytes): text to append to the file + become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not + provide a value, it will default to False. Defaults None + ignore_errors (bool, optional): Ignore errors. Defaults to False. + + Returns: + dict: Command result + + Raises: + Exception: If the file cannot be modified + """ + if self.get_host_variables(host)['os_name'] == 'windows': + file_content = self.get_file_content(host, path) + content = file_content + '\n' + block + result = self.modify_file_content(host, path, content) + + else: + become = self.get_host_variables(host).get('become', False) if become is None else become + + testinfra_host = self.get_host(host) + ansible_arguments = f"path='{path}' block='{block}'" + ansible_command = 'blockinfile' + + result = testinfra_host.ansible(ansible_command, ansible_arguments, check=False, become=become) + + if not result.get('msg', 'Block inserted') and not ignore_errors: + 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) \ No newline at end of file diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index ca3da66..96b8d05 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -9,7 +9,7 @@ DEFAULT_INSTALL_PATH = { 'linux': '/var/ossec', - 'windows': 'C:\\Program Files\\ossec-agent', + 'windows': 'C:\\Program Files (x86)\\ossec-agent', 'darwin': '/Library/Ossec' } @@ -51,7 +51,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 if os_host == 'windows' else os.path.join(installation_path, 'logs') + return installation_path + '\\logs' if os_host == 'windows' else os.path.join(installation_path, 'logs') def get_shared_directory_path(custom_installation_path=None, os_host='linux'): @@ -485,15 +485,6 @@ def get_ansible_host_component(self, host): manager_list = self.get_managers() return 'agent' if host in agent_list else 'manager' if host in manager_list else None - def restart_agent(self, host): - """Restart agent - - Args: - host (str): Hostname - systemd (bool, optional): Restart using systemd. Defaults to False. - """ - pass - def get_agents_info(self): """Get registered agents information. @@ -502,21 +493,31 @@ def get_agents_info(self): """ pass - def get_agents_id(self, agents_list=None): - """Get agents id - - Returns: - List: Agents id list + def get_agents_id(self, agent_list): + """Get agent ids + Args: + agents_list (_type_, agents_list): Agents list. + Return: + dict: agent_ids """ - pass + # Getting hostnames + host_names = [] + for agent in agent_list: + host_names.append(self.run_command(agent, 'hostname')[1]) - def restart_agents(self, agent_list=None, parallel=True): - """Restart list of agents + # Getting id - hostnames from manager + agent_control = self.run_command('manager1', '/var/ossec/bin/manage_agents -l', True)[1] - Args: - agent_list (list, optional): Agent list. Defaults to None. - """ - pass + # 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 @@ -576,6 +577,22 @@ def clean_client_keys(self, hosts=None): """ pass + def clean_logs(self, hosts): + """Remove host logs + Args: + hosts (_type_, hosts): host list. + """ + # 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) + def clean_agents(self, agents=None): """Stop agents, remove them from manager and clean their client keys @@ -584,18 +601,64 @@ def clean_agents(self, agents=None): """ pass - def remove_agents_from_manager(self, agents=None, status='all', older_than='0s'): + def restart_agents(self, agent_list): + """Restart agents + Args: + agents_list (_type_, agents_list): Agents list. + """ + # 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) + + + def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', parallel=True , logs=False, + restart=False): """Remove agents from manager Args: - agents (list, optional): Agents list. Defaults to None. - status (str, optional): Agents status. Defaults to 'all'. - older_than (str, optional): Older than parameter. Defaults to '0s'. - - Returns: - dict: API response - """ - pass + agent_list (list, optional): Agents list. Defaults to None. + 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. + restart (str): Restart agents. Defaults to False. + """ + if manager is None: manager = 'manager1' + if method == 'api': parallel = False + + # Getting agent_ids list + agent_ids = self.get_agents_id(agent_list) + + # Remove agent by cmd core function + def remove_agent_cmd(id): + print(id) + self.run_command(manager , f"/var/ossec/bin/manage_agents -r {id}", True) + + # Remove processes + if parallel: + if method == 'cmd': + self.pool.map(remove_agent_cmd, agent_ids) + else: + if method == 'cmd': + for id in agent_ids: + remove_agent_cmd(id) + elif method == 'api': + 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) + + # Remove logs + if logs: + self.clean_logs(agent_list) + + # Restarting agents + if restart: + self.restart_agents(agent_list) def stop_manager(self, manager): """Stop manager From da86fb57f72c558f523b43b85472c2ca9af3c008 Mon Sep 17 00:00:00 2001 From: Antonio Date: Wed, 7 Jun 2023 14:31:02 +0200 Subject: [PATCH 2/9] feat(#46): Method tested --- src/wazuh_qa_framework/system/host_manager.py | 2 +- .../system/wazuh_handler.py | 31 ++++++++++--------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index d324504..b5e0b8e 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -1343,4 +1343,4 @@ def make_api_call(self, host, port=55000, method='GET', endpoint='/', request_bo 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) \ No newline at end of file + 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 96b8d05..3a919a4 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -4,6 +4,7 @@ import os import re +from multiprocessing.pool import ThreadPool from wazuh_qa_framework.system.host_manager import HostManager @@ -139,8 +140,9 @@ def get_wazuh_file_path(custom_installation_path=None, os_host='linux', file_nam class WazuhEnvironmentHandler(HostManager): - def __init__(self, inventory_path): + def __init__(self, inventory_path, debug=False, max_workers=10): super().__init__(inventory_path) + self.pool = ThreadPool(max_workers) def get_file_fullpath(self, host, filename, group=None): """Get the path of common configuration and log file in the specified host. @@ -517,6 +519,7 @@ def get_agents_id(self, agent_list): id_value = line.split(',')[0].split(': ')[1].strip() agent_ids.append(id_value) break + return agent_ids def restart_manager(self, host): @@ -592,7 +595,8 @@ def clean_logs(self, hosts): 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) - + + def clean_agents(self, agents=None): """Stop agents, remove them from manager and clean their client keys @@ -601,6 +605,7 @@ def clean_agents(self, agents=None): """ pass + def restart_agents(self, agent_list): """Restart agents Args: @@ -614,7 +619,7 @@ def restart_agents(self, agent_list): 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, restart=False): """Remove agents from manager @@ -635,22 +640,20 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par # Remove agent by cmd core function def remove_agent_cmd(id): - print(id) self.run_command(manager , f"/var/ossec/bin/manage_agents -r {id}", True) - + # Remove processes - if parallel: - if method == 'cmd': + if method == 'cmd': + if parallel: self.pool.map(remove_agent_cmd, agent_ids) - else: - if method == 'cmd': + else: for id in agent_ids: remove_agent_cmd(id) - elif method == 'api': - 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) + 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) # Remove logs if logs: From e694057708a8d6337bf980a7cb318e4e58879e74 Mon Sep 17 00:00:00 2001 From: Antonio Date: Wed, 7 Jun 2023 15:10:12 +0200 Subject: [PATCH 3/9] feat(#46): Changes after linter --- src/wazuh_qa_framework/system/host_manager.py | 643 ------------------ .../system/wazuh_handler.py | 11 +- 2 files changed, 5 insertions(+), 649 deletions(-) diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index b5e0b8e..3ed42b8 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -2,649 +2,6 @@ # Created by Wazuh, Inc. . # This program is free software; you can redistribute it and/or modify it under the terms of GPLv2 -import tempfile -import testinfra -import base64 -import os -import json -from ansible.inventory.manager import InventoryManager -from ansible.parsing.dataloader import DataLoader -from ansible.vars.manager import VariableManager - - -class HostManager: - """Remote host management interface. - - It allows to manage remote hosts using ansible inventory and testinfra framework. - - Args: - inventory_path (str): Ansible inventory path - - Attributes: - inventory_path (str): Ansible inventory path - inventory_manager (ansible.inventory.manager.InventoryManager): Ansible inventory manager - variable_manager (ansible.vars.manager.VariableManager): Ansible variable manager - """ - - def __init__(self, inventory_path): - self.inventory_path = inventory_path - - data_loader = DataLoader() - self.inventory_manager = InventoryManager(loader=data_loader, sources=inventory_path) - self.hosts_variables = {} - - variable_manager = VariableManager(loader=data_loader, inventory=self.inventory_manager) - - for host in self.inventory_manager.get_hosts(): - self.hosts_variables[host] = variable_manager.get_vars(host=self.inventory_manager.get_host(str(host))) - - def get_host(self, host): - """Get the testinfra host. - - Args: - host (str): Hostname - - Returns: - testinfra.modules.base.Ansible: Host instance from hostspec - """ - return testinfra.get_host(f"ansible://{host}?ansible_inventory={self.inventory_path}") - - def get_groups(self): - """Get the groups of the inventory. - - Returns: - list: Groups of the inventory - """ - return list(self.inventory_manager.groups.keys()) - - def get_group_hosts(self, pattern=None): - """Get all hosts from inventory that belong to a group. - - Args: - group (str): Group name - - Returns: - list: List of hosts - """ - if pattern: - return [str(host) for host in self.inventory_manager.get_hosts(pattern=pattern)] - else: - return [str(host) for host in self.inventory_manager.get_hosts()] - - def get_host_variables(self, host): - """Get the variables of the specified host. - - Args: - host (str): Hostname - - Returns: - testinfra.modules.base.Ansible: Host instance from hostspec - """ - inventory_manager_host = self.inventory_manager.get_host(host) - - return self.hosts_variables[inventory_manager_host] - - def collect_host_ansible_facts(self, host): - """Get the ansible facts of the specified host. - - Args: - host (str): Hostname - - Returns: - str: OS of the host - """ - testinfra_host = self.get_host(host) - - return testinfra_host.ansible("setup") - - def collect_host_os(self, host): - """Get the OS of the specified host. - - Args: - host: Hostname - - Returns: - tuple: Hostname, Major version, Distribution version. Example: ('CentOS', '7', '7.6.1810') - """ - ansible_facts = self.collect_host_ansible_facts(host) - - return (ansible_facts['ansible_facts']['ansible_distribution'], - ansible_facts['ansible_facts']['ansible_distribution_major_version'], - ansible_facts['ansible_facts']['ansible_distribution_version']) - - def collect_host_ips(self, host): - """Get the host IPs - - Args: - host (str): Hostname - - Returns: - dict: IPs of the host (ipv4 and ipv6). Example: {'ipv4': ['172.31.5.209'], 'ipv6': ['fe80::f::fef4:bb6d']} - """ - ansible_facts = self.collect_host_ansible_facts(host) - - return {'ipv4': ansible_facts['ansible_facts']['ansible_all_ipv4_addresses'], - 'ipv6': ansible_facts['ansible_facts']['ansible_all_ipv6_addresses']} - - def collect_host_interfaces(self, host): - """Get the interfaces of the specified host. - - Args: - host (str): Hostname - - Returns: - list: Interfaces of the host. Example ['lo', 'eth0'] - """ - ansible_facts = self.collect_host_ansible_facts(host) - - return ansible_facts['ansible_facts']['ansible_interfaces'] - - def check_connection(self, host): - """Check if the host is reachable. - - Args: - host (str): Hostname - - Returns: - bool: True if the host is reachable, False otherwise - """ - testinfra_host = self.get_host(host) - ansible_command = 'win_ping' if self.get_host_variables(host)['os_name'] == 'windows' else 'ping' - return testinfra_host.ansible(ansible_command, check=False)['ping'] == 'pong' - - def copy_file(self, host, src_path, dest_path, remote_src=False, become=None, ignore_errors=False): - """Move from src_path to the desired location dest_path for the specified host. - - Args: - host (str): Hostname - src_path (str): Source path - dest_path (str): Destination path - remote_src (bool): If True, the file is assumed to live on the remote machine, not the controller. - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool): Ignore errors - - Returns: - dict: Result of the command execution - - Raises: - Exception: If the command execution fails - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - ansible_command = 'win_copy' if self.get_host_variables(host)['os_name'] == 'windows' else 'copy' - remote_source = 'yes' if remote_src else 'no' - - command_parameters = f"src='{src_path}' dest='{dest_path}' remote_src={remote_source}" - result = testinfra_host.ansible(ansible_command, command_parameters, check=False, become=become) - - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error moving file from {src_path} to {dest_path} on host {host}: {result}") - - return result - - def get_file_content(self, host, path, become=None, ignore_errors=False): - """Read a file from the specified host. - - Args: - host (str): Hostname - path (str): File path - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool): Ignore errors - - Returns: - str: File content - - Raises: - Exception: If the file cannot be read - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - result = testinfra_host.ansible("slurp", f"src='{path}'", check=False, become=become) - - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error reading file {path} on host {host}: {result}") - - return base64.b64decode(result['content']).decode('utf-8') - - def synchronize_linux_directory(self, host, dest_path, src_path=None, filesystem=None, become=None, - ignore_errors=False): - """Create a file structure on the specified host. - Not supported on Windows. - - Args: - host (str): Hostname - dest_path (str): Destination path - filesystem (dict): File structure - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Result of the command execution - - Raises: - Exception: If the command execution fails - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - - ansible_command = 'synchronize' - - if filesystem: - tmp_directory = tempfile.TemporaryDirectory() - directory_path = os.path.join(tmp_directory.name, filesystem['directory_name']) - os.mkdir(directory_path) - src_path = directory_path - - for file in filesystem['files']: - file_path = f"{directory_path}/{file['filename']}" - with open(file_path, 'w') as file_operator: - file_operator.write(file['content']) - - result = testinfra_host.ansible(ansible_command, f"src='{src_path}' dest='{dest_path}'", check=False, - become=become) - - if (result['rc'] != 0 or not result) and not ignore_errors: - raise Exception(f"Error creating file structure on host {host}: {result}") - - return result - - def truncate_file(self, host, file_path, recreate=True, become=None, ignore_errors=False): - """Truncate a file from the specified host. - - Args: - host (str): Hostname - file_path (str): File path - recreate (bool, optional): Recreate file. Defaults to True. - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the file cannot be truncated - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - result = None - - if recreate: - ansible_command = 'win_copy' if self.get_host_variables(host)['os_name'] == 'windows' else 'copy' - result = testinfra_host.ansible(ansible_command, f"dest='{file_path}' content=''", check=False, - become=become) - else: - ansible_command = 'win_file' if self.get_host_variables(host)['os_name'] == 'windows' else 'file' - result = testinfra_host.ansible(ansible_command, f"path='{file_path}' state=touch", check=False, - become=become) - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error truncating file {file_path} on host {host}: {result}") - - return result - - def remove_file(self, host, file_path, become=None, ignore_errors=False): - """Remove a file from the specified host. - - Args: - host (str): Hostname - file_path (str): File path - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the file cannot be removed - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - ansible_command = 'win_file' if self.get_host_variables(host)['os_name'] == 'windows' else 'file' - result = testinfra_host.ansible(ansible_command, f"path='{file_path}' state=absent", check=False, become=become) - - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error removing file {file_path} on host {host}: {result}") - - return result - - def modify_file_content(self, host, path, content, become=None, ignore_errors=False): - """Create a file with a specified content and copies it to a path. - - Args: - host (str): Hostname - path (str): path for the file to create and modify - content (str, bytes): content to write into the file - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the file cannot be modified - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - tmp_file = tempfile.NamedTemporaryFile() - with open(tmp_file.name, 'w+') as tmp: - tmp.write(content) - - result = self.copy_file(host, src_path=tmp_file.name, dest_path=path, become=become) - - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error modifying file {path} on host {host}: {result}") - - return result - - def create_file(self, host, path, content, directory=False, owner=None, group=None, mode=None, become=None, - ignore_errors=False): - """Create a file with a specified content and copies it to a path. - - Args: - host (str): Hostname - path (str): path for the file to create and modify - content (str, bytes): content to write into the file - owner (str): owner of the file - group (str): group of the file - mode (str): mode of the file - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the file cannot be created - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - tmp_file = tempfile.NamedTemporaryFile() - with open(tmp_file.name, 'w+') as tmp: - tmp.write(content) - - ansible_command = 'win_copy' if self.get_host_variables(host)['os_name'] == 'windows' else 'copy' - - ansible_parameters = f"src='{tmp_file.name}' dest='{path}'" - ansible_parameters += f" owner={owner}" if owner else '' - ansible_parameters += f" group={group}" if group else '' - ansible_parameters += f" mode={mode}" if mode else '' - ansible_parameters += f' state=directory' if directory else '' - - result = testinfra_host.ansible(ansible_command, ansible_parameters, check=False, become=become) - - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error creating file {path} on host {host}: {result}") - return result - - def control_service(self, host, service, state, become=None, ignore_errors=False): - """Control a service on a host. - - Args: - host (str): Hostname - service (str): Service name - state (str): Service state - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the service cannot be controlled - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - ansible_command = 'win_service' if self.get_host_variables(host)['os_name'] == 'windows' else 'service' - - result = testinfra_host.ansible(ansible_command, f"name={service} state={state}", check=False, become=become) - - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error controlling service {service} on host {host}: {result}") - - return result - - def run_command(self, host, cmd, become=None, ignore_errors=False): - """Run a command on a host. - - Args: - host (str): Hostname - cmd (str): Command to run - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the command cannot be run - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - ansible_command = 'win_command' if self.get_host_variables(host)['os_name'] == 'windows' else 'command' - - result = testinfra_host.ansible(ansible_command, f"{cmd}", check=False, become=become) - rc, stdout = result.get('rc', 1), result.get('stdout', '') - - if rc != 0 and not ignore_errors: - raise Exception(f"Error running command '{cmd}' on host {host}: {result}") - - return rc, stdout - - def run_shell(self, host, cmd, become=None, ignore_errors=False): - """Run a shell command on a host. - The difference with run_command is that here, shell symbols like &, |, etc. are interpreted. - - Args: - host (str): Hostname - cmd (str): Command to run - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the command cannot be run - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - rc = None - stdout = None - - ansible_command = 'win_shell' if self.get_host_variables(host)['os_name'] == 'windows' else 'shell' - - result = testinfra_host.ansible(ansible_command, f"{cmd}", check=False, become=become) - - rc, stdout = result['rc'], result['stdout'] - - if rc != 0 and not ignore_errors: - raise Exception(f"Error running command {cmd} on host {host}: {result}") - - return rc, stdout - - def find_files(self, host, path, pattern, recurse=False, use_regex=False, become=None, ignore_errors=False): - """Search and return information of a file inside a path. - - Args: - host (str): Hostname - path (str): Path in which to search for the file that matches the pattern. - pattern (str): Restrict the files to be returned to those whose basenames match the pattern specified. - recurse (bool): If target is a directory, recursively descend into the directory looking for files. - use_regex (bool): If no, the patterns are file globs (shell), if yes, they are python regexes. - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - Files (list): List of found files. - - Raises: - Exception: If the command cannot be run - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - test_infra_host = self.get_host(host) - ansible_command = 'win_find' if self.get_host_variables(host)['os_name'] == 'windows' else 'find' - ansible_pattern_arguments = 'patterns' if self.get_host_variables(host)['os_name'] == 'windows' else 'pattern' - - result = test_infra_host.ansible(ansible_command, f"paths={path} {ansible_pattern_arguments}='{pattern}' \ - recurse={recurse} use_regex={use_regex}", - become=become, check=False) - - if 'files' not in result and not ignore_errors: - raise Exception(f"Error finding file {path} on host {host}: {result}") - - return result['files'] - - def get_file_stats(self, host, path, become=None, ignore_errors=False): - """Retrieve file or file system status. - - Args: - host (str): Hostname. - path (str): The full path of the file/object to get the facts of. - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result. - - Raises: - Exception: If the command cannot be run. - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - - if self.get_host_variables(host)['os_name'] == 'windows': - ansible_command = 'ansible.windows.win_stat' - else: - ansible_command = 'stat' - - result = testinfra_host.ansible(ansible_command, f"path='{path}'", check=False, become=become) - - if 'stat' not in result and not ignore_errors: - raise Exception(f"Error getting stats of {path} on host {host}: {result}") - - return result - - def install_package(self, host, package_name, become=None, ignore_errors=False): - """Install a package on a host. - - Args: - host (str): Hostname - package_name (str): Package to install - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the install cannot be run - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - ansible_arguments = f"name={package_name} state=present" - ansible_command = 'win_chocolatey' if self.get_host_variables(host)['os_name'] == 'windows' else 'package' - - result = testinfra_host.ansible(ansible_command, ansible_arguments, - check=False, become=become) - - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error installing package {package_name} on host {host}: {result}") - - return result - - def uninstall_package(self, host, package_name, become=None, ignore_errors=False): - """Uninstall a package on a host. - - Args: - host (str): Hostname - package_name (str): Package to uninstall - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the uninstall cannot be run - """ - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - ansible_arguments = f"name={package_name} state=absent" - ansible_command = 'win_chocolatey' if self.get_host_variables(host)['os_name'] == 'windows' else 'package' - - result = testinfra_host.ansible(ansible_command, ansible_arguments, - check=False, become=become) - - if result.get('msg', None) and not ignore_errors: - raise Exception(f"Error installing package {package_name} on host {host}: {result}") - - return result - - def append_block_in_file(self, host, path, block, become=None, ignore_errors=False): - """Append a text block in file - - Args: - host (str): Hostname - path (str): path for the file to insert a block - block (str, bytes): text to append to the file - become (bool): If no value is provided, it will be taken from the inventory. If the inventory does not - provide a value, it will default to False. Defaults None - ignore_errors (bool, optional): Ignore errors. Defaults to False. - - Returns: - dict: Command result - - Raises: - Exception: If the file cannot be modified - """ - if self.get_host_variables(host)['os_name'] == 'windows': - file_content = self.get_file_content(host, path) - content = file_content + '\n' + block - result = self.modify_file_content(host, path, content) - - else: - become = self.get_host_variables(host).get('become', False) if become is None else become - - testinfra_host = self.get_host(host) - ansible_arguments = f"path='{path}' block='{block}'" - ansible_command = 'blockinfile' - - result = testinfra_host.ansible(ansible_command, ansible_arguments, check=False, become=become) - - if not result.get('msg', 'Block inserted') and not ignore_errors: - raise Exception(f"Error inserting a block in file {path} on host {host}: {result}") - - return result -# Copyright (C) 2015-2021, Wazuh Inc. -# Created by Wazuh, Inc. . -# This program is free software; you can redistribute it and/or modify it under the terms of GPLv2 - import tempfile import testinfra import base64 diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index 3a919a4..8325c82 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -123,7 +123,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'], @@ -589,17 +589,16 @@ def clean_logs(self, hosts): 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) + 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) + 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) + 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 - Args: agents (_type_, agents_list): Agents list. Defaults to None. """ @@ -640,7 +639,7 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par # Remove agent by cmd core function def remove_agent_cmd(id): - self.run_command(manager , f"/var/ossec/bin/manage_agents -r {id}", True) + self.run_command(manager, f"/var/ossec/bin/manage_agents -r {id}", True) # Remove processes if method == 'cmd': From f5c1621fc4e18091b1837f3bb7561710ad38d002 Mon Sep 17 00:00:00 2001 From: Antonio Date: Wed, 7 Jun 2023 15:20:05 +0200 Subject: [PATCH 4/9] feat(#46): import json --- src/wazuh_qa_framework/system/host_manager.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/wazuh_qa_framework/system/host_manager.py b/src/wazuh_qa_framework/system/host_manager.py index 3ed42b8..ee30a64 100644 --- a/src/wazuh_qa_framework/system/host_manager.py +++ b/src/wazuh_qa_framework/system/host_manager.py @@ -6,6 +6,7 @@ import testinfra import base64 import os +import json from ansible.inventory.manager import InventoryManager from ansible.parsing.dataloader import DataLoader from ansible.vars.manager import VariableManager From b15ad7e4c9f43573259ca0d1f9a0b806acd118a2 Mon Sep 17 00:00:00 2001 From: Antonio Date: Fri, 9 Jun 2023 12:40:05 +0200 Subject: [PATCH 5/9] feat(#46): More changes after linter --- src/wazuh_qa_framework/system/wazuh_handler.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index 8325c82..773f7d9 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -596,7 +596,6 @@ def clean_logs(self, hosts): if 'master' == host_type or 'worker' == host_type: 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 Args: @@ -604,7 +603,6 @@ def clean_agents(self, agents=None): """ pass - def restart_agents(self, agent_list): """Restart agents Args: @@ -618,8 +616,7 @@ def restart_agents(self, agent_list): 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 @@ -631,8 +628,8 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par logs (str): Remove logs from agents. Defaults to False. restart (str): Restart agents. Defaults to False. """ - if manager is None: manager = 'manager1' - if method == 'api': parallel = False + manager = 'manager1' if manager is None else manager + parallel = False if method == 'api' else parallel # Getting agent_ids list agent_ids = self.get_agents_id(agent_list) @@ -650,8 +647,8 @@ def remove_agent_cmd(id): remove_agent_cmd(id) 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', + 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) # Remove logs From 172235f514b1ae066d0c64415c936eb2b8ba1222 Mon Sep 17 00:00:00 2001 From: Antonio Date: Fri, 9 Jun 2023 12:41:27 +0200 Subject: [PATCH 6/9] feat(#46): New changes after linter --- src/wazuh_qa_framework/system/wazuh_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index 773f7d9..be59a2c 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -616,7 +616,7 @@ def restart_agents(self, agent_list): 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 From 508698d1b664d38a5b033d782ae0a839951eb646 Mon Sep 17 00:00:00 2001 From: Antonio Date: Thu, 15 Jun 2023 11:18:53 +0200 Subject: [PATCH 7/9] 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 From 75bb26a35b44e7722d3b15d854d33f22aeacb758 Mon Sep 17 00:00:00 2001 From: Antonio Date: Thu, 15 Jun 2023 16:08:07 +0200 Subject: [PATCH 8/9] 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 From cc9476bf94768042987e5c169607e514362b6a47 Mon Sep 17 00:00:00 2001 From: Antonio Date: Thu, 15 Jun 2023 16:19:08 +0200 Subject: [PATCH 9/9] fix(#46): Fixes after linter --- src/wazuh_qa_framework/system/wazuh_handler.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/wazuh_qa_framework/system/wazuh_handler.py b/src/wazuh_qa_framework/system/wazuh_handler.py index a406552..44094ab 100644 --- a/src/wazuh_qa_framework/system/wazuh_handler.py +++ b/src/wazuh_qa_framework/system/wazuh_handler.py @@ -513,7 +513,7 @@ def get_agent_id(self, host, agent): Return: str: agent_id """ - host_list = WazuhAPI(address = self.get_host_variables(host)['ip']).list_agents()['affected_items'] + 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') @@ -830,7 +830,8 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par logs (str): Remove logs (ossec.log, api.log) from agents. Defaults to False. restart (str): Restart agents. Defaults to False. """ - if manager is None: manager = 'manager1' + if manager is None: + manager = 'manager1' # Getting agent_ids list agent_ids = [] @@ -849,15 +850,15 @@ def remove_agents_from_manager(self, agent_list, manager=None, method='cmd', par 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' + 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'])) + request.send(WazuhAPI(address=self.get_host_variables(manager)['ip'])) # Remove logs - if logs and parallel == False: + if logs and not parallel: for agent in agent_list: self.clean_logs(agent) - if logs and parallel == True: + if logs and parallel: self.pool.map(self.clean_logs, agent_list) # Restarting agents