Implementation of agent remove method in WazuhHandler

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,64 @@ 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)

src/wazuh_qa_framework/system/wazuh_handler.py

@@ -4,12 +4,13 @@
 
 import os
 import re
+from multiprocessing.pool import ThreadPool
 
 from wazuh_qa_framework.system.host_manager import HostManager
 
 DEFAULT_INSTALL_PATH = {
     'linux': '/var/ossec',
-    'windows': 'C:\\Program Files\\ossec-agent',
+    'windows': 'C:\\Program Files (x86)\\ossec-agent',
     'darwin': '/Library/Ossec'
 }
 
@@ -51,7 +52,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'):
@@ -122,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'],
@@ -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.
@@ -485,15 +487,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 +495,32 @@ 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,26 +580,84 @@ 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
-
         Args:
             agents (_type_, agents_list): Agents list. Defaults to 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
+            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.
         """
-        pass
+        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)
+
+        # Remove agent by cmd core function
+        def remove_agent_cmd(id):
+            self.run_command(manager, f"/var/ossec/bin/manage_agents -r {id}", True)
+
+        # Remove processes
+        if method == 'cmd':
+            if parallel:
+                self.pool.map(remove_agent_cmd, agent_ids)
+            else:
+                for id in agent_ids:
+                    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',
+                               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