enhancement(#5229): Improving get_os_type and arch

wazuh · Apr 19, 2024 · 684ea5d · 684ea5d
1 parent 8aa767e
commit 684ea5d
Show file tree

Hide file tree

Showing 10 changed files with 124 additions and 31 deletions.
diff --git a/deployability/modules/testing/tests/helpers/agent.py b/deployability/modules/testing/tests/helpers/agent.py
@@ -8,7 +8,7 @@
 from .constants import WAZUH_CONF, WAZUH_ROOT
 from .executor import Executor, WazuhAPI
 from .generic import HostInformation, CheckFiles
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 
 class WazuhAgent:
 
@@ -29,19 +29,19 @@ def install_agent(inventory_path, agent_name, wazuh_version, wazuh_revision, liv
             distribution = HostInformation.get_linux_distribution(inventory_path)
             architecture = HostInformation.get_architecture(inventory_path)
 
-            if distribution == 'rpm' and 'x86_64' in architecture:
+            if distribution == 'rpm' and 'amd64' in architecture:
                 commands.extend([
                     f"curl -o wazuh-agent-{wazuh_version}-1.x86_64.rpm https://{s3_url}.wazuh.com/{release}/yum/wazuh-agent-{wazuh_version}-1.x86_64.rpm && sudo WAZUH_MANAGER='MANAGER_IP' WAZUH_AGENT_NAME='{agent_name}' rpm -ihv wazuh-agent-{wazuh_version}-1.x86_64.rpm"
                 ])
-            elif distribution == 'rpm' and 'aarch64' in architecture:
+            elif distribution == 'rpm' and 'arm64' in architecture:
                 commands.extend([
                     f"curl -o wazuh-agent-{wazuh_version}-1aarch64.rpm https://{s3_url}.wazuh.com/{release}/yum/wazuh-agent-{wazuh_version}-1.aarch64.rpm && sudo WAZUH_MANAGER='MANAGER_IP' WAZUH_AGENT_NAME='{agent_name}' rpm -ihv wazuh-agent-{wazuh_version}-1.aarch64.rpm"
                 ])
-            elif distribution == 'deb' and 'x86_64' in architecture:
+            elif distribution == 'deb' and 'amd64' in architecture:
                 commands.extend([
                     f"wget https://{s3_url}.wazuh.com/{release}/apt/pool/main/w/wazuh-agent/wazuh-agent_{wazuh_version}-1_amd64.deb && sudo WAZUH_MANAGER='MANAGER_IP' WAZUH_AGENT_NAME='{agent_name}' dpkg -i ./wazuh-agent_{wazuh_version}-1_amd64.deb"
                 ])
-            elif distribution == 'deb' and 'aarch64' in architecture:
+            elif distribution == 'deb' and 'arm64' in architecture:
                 commands.extend([
                     f"wget https://{s3_url}.wazuh.com/{release}/apt/pool/main/w/wazuh-agent/wazuh-agent_{wazuh_version}-1_arm64.deb && sudo WAZUH_MANAGER='MANAGER_IP' WAZUH_AGENT_NAME='{agent_name}' dpkg -i ./wazuh-agent_{wazuh_version}-1arm64.deb"
                 ])

diff --git a/deployability/modules/testing/tests/helpers/generic.py b/deployability/modules/testing/tests/helpers/generic.py
@@ -14,7 +14,7 @@
 from pathlib import Path
 from .constants import WAZUH_CONTROL, CLIENT_KEYS
 from .executor import Executor
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 from .utils import Utils
 
 
@@ -69,16 +69,19 @@ def get_os_type(inventory_path) -> str:
         Returns:
             str: type of host (windows, linux, macos)
         """
-        if 'manager' in inventory_path:
-            pattern = r'manager-(\w+)-'
-        elif 'agent' in inventory_path:
-            pattern = r'agent-(\w+)-'
-        result = re.search(pattern, inventory_path)
-        if result:
-            return result.group(1)
-        else:
-            return None
-
+        try:
+            with open(inventory_path.replace('inventory', 'track'), 'r') as file:
+                data = yaml.safe_load(file)
+            if 'platform' in data:
+                return data['platform']
+            else:
+                raise KeyError("The 'platform' key was not found in the YAML file.")
+        except FileNotFoundError:
+            logger.error(f"The YAML file '{inventory_path}' was not found.")
+        except yaml.YAMLError as e:
+            logger.error(f"Error while loading the YAML file: {e}")
+        except Exception as e:
+            logger.error(f"An unexpected error occurred: {e}")
 
 
     @staticmethod
@@ -90,9 +93,21 @@ def get_architecture(inventory_path) -> str:
             inventory_path: host's inventory path
 
         Returns:
-            str: architecture (aarch64, x86_64, intel, apple)
+            str: architecture (amd64, arm64, intel, apple)
         """
-        return Executor.execute_command(inventory_path, 'uname -m')
+        try:
+            with open(inventory_path.replace('inventory', 'track'), 'r') as file:
+                data = yaml.safe_load(file)
+            if 'platform' in data:
+                return data['arch']
+            else:
+                raise KeyError("The 'platform' key was not found in the YAML file.")
+        except FileNotFoundError:
+            logger.error(f"The YAML file '{inventory_path}' was not found.")
+        except yaml.YAMLError as e:
+            logger.error(f"Error while loading the YAML file: {e}")
+        except Exception as e:
+            logger.error(f"An unexpected error occurred: {e}")
 
 
     @staticmethod
@@ -692,7 +707,7 @@ def hasAgentClientKeys(inventory_path) -> bool:
         """
         os_type = HostInformation.get_os_type(inventory_path)
         if os_type == 'linux':
-            return HostInformation.file_exists(inventory_path,{CLIENT_KEYS})
+            return HostInformation.file_exists(inventory_path, {CLIENT_KEYS})
         elif os_type == 'macos':
             return HostInformation.file_exists(inventory_path, '/Library/Ossec/etc/client.keys')
 

diff --git a/deployability/modules/testing/tests/helpers/manager.py b/deployability/modules/testing/tests/helpers/manager.py
@@ -8,7 +8,7 @@
 from .constants import CLUSTER_CONTROL, AGENT_CONTROL, WAZUH_CONF, WAZUH_ROOT
 from .executor import Executor, WazuhAPI
 from .generic import HostInformation, CheckFiles
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 from .utils import Utils
 
 

diff --git a/deployability/modules/testing/tests/helpers/utils.py b/deployability/modules/testing/tests/helpers/utils.py
@@ -8,7 +8,7 @@
 import logging
 import time
 
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 
 
 paramiko_logger = logging.getLogger("paramiko")

diff --git a/deployability/modules/testing/tests/test_agent/test_install.py b/deployability/modules/testing/tests/test_agent/test_install.py
@@ -8,7 +8,7 @@
 from ..helpers.agent import WazuhAgent
 from ..helpers.constants import WAZUH_ROOT
 from ..helpers.generic import HostConfiguration, HostInformation, GeneralComponentActions
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 from ..helpers.manager import WazuhManager
 from ..helpers.utils import Utils
 

diff --git a/deployability/modules/testing/tests/test_agent/test_registration.py b/deployability/modules/testing/tests/test_agent/test_registration.py
@@ -8,7 +8,7 @@
 from ..helpers.agent import WazuhAgent, WazuhAPI
 from ..helpers.generic import HostInformation, GeneralComponentActions, Waits
 from ..helpers.manager import WazuhManager, WazuhAPI
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 from ..helpers.utils import Utils
 
 @pytest.fixture(scope="module", autouse=True)

diff --git a/deployability/modules/testing/tests/test_agent/test_restart.py b/deployability/modules/testing/tests/test_agent/test_restart.py
@@ -6,7 +6,7 @@
 import re
 
 from ..helpers.generic import GeneralComponentActions, HostInformation
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 from ..helpers.manager import WazuhManager
 from ..helpers.utils import Utils
 

diff --git a/deployability/modules/testing/tests/test_agent/test_stop.py b/deployability/modules/testing/tests/test_agent/test_stop.py
@@ -7,7 +7,7 @@
 
 from ..helpers.agent import WazuhAgent, WazuhAPI
 from ..helpers.generic import GeneralComponentActions, Waits, HostInformation
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 from ..helpers.utils import Utils
 
 @pytest.fixture(scope="module", autouse=True)

diff --git a/deployability/modules/testing/tests/test_agent/test_uninstall.py b/deployability/modules/testing/tests/test_agent/test_uninstall.py
@@ -9,7 +9,7 @@
 from ..helpers.constants import WAZUH_ROOT
 from ..helpers.generic import HostInformation, GeneralComponentActions, Waits
 from ..helpers.manager import WazuhManager, WazuhAPI
-from modules.generic.logger import logger
+from modules.testing.utils import logger
 from ..helpers.utils import Utils
 
 @pytest.fixture(scope="module", autouse=True)

diff --git a/deployability/modules/workflow_engine/examples/agent/vagrant/test.yaml b/deployability/modules/workflow_engine/examples/agent/vagrant/test.yaml
@@ -1,14 +1,91 @@
 version: 0.1
-description: Test agent restart with provisioning agents' with provision module
+description: This workflow is used to test agents' deployment for DDT1 PoC
 variables:
   agent-os:
-    -  macos-ventura-13.4.1-arm64
-  manager-os: linux-ubuntu-18.04-amd64
-  infra-provider: aws
+    - linux-ubuntu-18.04-amd64
+
+  manager-os: linux-ubuntu-22.04-amd64
+  infra-provider: vagrant
   working-dir: /tmp/dtt1-poc
 
 tasks:
+  # Unique manager allocate task
+  - task: "allocate-manager-{manager-os}"
+    description: "Allocate resources for the manager."
+    do:
+      this: process
+      with:
+        path: python3
+        args:
+          - modules/allocation/main.py
+          - action: create
+          - provider: "{infra-provider}"
+          - size: large
+          - composite-name: "{manager-os}"
+          - inventory-output: "{working-dir}/manager-{manager-os}/inventory.yaml"
+          - track-output: "{working-dir}/manager-{manager-os}/track.yaml"
+          - label-termination-date: "1d"
+          - label-team: "qa"
+    on-error: "abort-all"
+    cleanup:
+      this: process
+      with:
+        path: python3
+        args:
+          - modules/allocation/main.py
+          - action: delete
+          - track-output: "{working-dir}/manager-{manager-os}/track.yaml"
 
+  # Unique agent allocate task
+  - task: "allocate-agent-{agent}"
+    description: "Allocate resources for the agent."
+    do:
+      this: process
+      with:
+        path: python3
+        args:
+          - modules/allocation/main.py
+          - action: create
+          - provider: "{infra-provider}"
+          - size: small
+          - composite-name: "{agent}"
+          - inventory-output: "{working-dir}/agent-{agent}/inventory.yaml"
+          - track-output: "{working-dir}/agent-{agent}/track.yaml"
+          - label-termination-date: "1d"
+          - label-team: "qa"
+    on-error: "abort-all"
+    foreach:
+      - variable: agent-os
+        as: agent
+    cleanup:
+      this: process
+      with:
+        path: python3
+        args:
+          - modules/allocation/main.py
+          - action: delete
+          - track-output: "{working-dir}/agent-{agent}/track.yaml"
+    depends-on:
+      - "provision-manager-{manager-os}"
+
+  # Unique manager provision task
+  - task: "provision-manager-{manager-os}"
+    description: "Provision the manager."
+    do:
+      this: process
+      with:
+        path: python3
+        args:
+          - modules/provision/main.py
+          - inventory: "{working-dir}/manager-{manager-os}/inventory.yaml"
+          - install:
+            - component: wazuh-manager
+              type: assistant
+              version: 4.7.3
+              live: True
+    depends-on:
+      - "allocate-manager-{manager-os}"
+    on-error: "abort-all"
 
 
   # Generic agent test task
@@ -31,4 +108,5 @@ tasks:
     foreach:
       - variable: agent-os
         as: agent
-
+    depends-on:
+      - "allocate-agent-{agent}"