Merge pull request #5138 from wazuh/enhancement/4901-dtt1-iteration-3…

…-allocator-validations-and-error-handling

Improvements in error handling and some validations

deployability/modules/allocation/aws/credentials.py

@@ -30,7 +30,7 @@ def __init__(self) -> None:
         super().__init__()
         self._resource = boto3.resource('ec2')
 
-    def generate(self, base_dir: str | Path, name: str, overwrite: bool = False) -> Path:
+    def generate(self, base_dir: str | Path, name: str) -> Path:
         """
         Generates a new key pair and returns it.
 
@@ -58,11 +58,7 @@ def generate(self, base_dir: str | Path, name: str, overwrite: bool = False) ->
             # Check if the key pair already exists
             key_pair = self._resource.KeyPair(name)
             if key_pair.key_pair_id:
-                if not overwrite:
                     raise self.CredentialsError(f"Key pair {name} already exists.")
-                else:
-                    logger.warning(f"Key pair {name} already exists. Overwriting.")
-                    key_pair.delete()
         except ClientError:
             pass
 

deployability/modules/allocation/aws/models.py

@@ -9,7 +9,7 @@ class AWSConfig(ProviderConfig):
     type: str
     security_groups: list[str]
     termination_date: str
-    issue: str
+    issue: str | None = None
     team: str
     name: str
     host_identifier: str | None = None

deployability/modules/allocation/aws/provider.py

@@ -2,8 +2,9 @@
 import fnmatch
 import os
 import re
-import sys
+import random
 from pathlib import Path
+from datetime import datetime, timedelta
 
 from modules.allocation.generic import Provider
 from modules.allocation.generic.models import CreationPayload, InstancePayload, InstancePayload
@@ -53,9 +54,13 @@ def _create_instance(cls, base_dir: Path, params: CreationPayload, config: AWSCo
             host_identifier = None
             date_regex = r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}"
             url_regex = "(https:\/\/|http:\/\/)?[github]{2,}(\.[com]{2,})?\/wazuh\/[a-zA-Z0-9_-]+(?:-[a-zA-Z0-9_-]+)?\/issues\/[0-9]{2,}"
-            if not termination_date or not re.match(date_regex, termination_date):
-                logger.error(f"The termination_date label was not provided or is of incorrect format, example: 2024-02-25 12:00:00.")
-                sys.exit(1)
+            if not termination_date:
+                raise ValueError(f"The termination_date label was not provided.")
+            elif re.match(r'^\d+d$', termination_date):
+                new_date = datetime.now() + timedelta(days=int(termination_date.split("d")[0]))
+                termination_date = new_date.strftime("%Y-%m-%d %H:%M:%S")
+            elif not re.match(date_regex, termination_date):
+                raise ValueError(f"The termination_date label was not provided or is of incorrect format, example: 2021-12-31 23:59:59 or 2d")
             if label_team:
                 not_match = 0
                 for team in teams:
@@ -65,27 +70,26 @@ def _create_instance(cls, base_dir: Path, params: CreationPayload, config: AWSCo
                     else:
                         not_match += 1
                 if not_match == len(teams):
-                    logger.error(f"The team label provided does not match any of the available teams.")
-                    sys.exit(1)
+                    raise ValueError(f"The team label provided does not match any of the available teams. Available teams: {teams}")
             else:
-                logger.error(f"The team label was not provided.")
-                sys.exit(1)
-            if not issue or not re.match(url_regex, issue):
-                logger.error(f"The issue label was not provided or is of incorrect format, example: https://github.com/wazuh/<repository>/issues/<issue-number>.")
-                sys.exit(1)
+                raise ValueError(f"The team label was not provided. Availables teams: {teams}.")
             if params.instance_name:
                 name = params.instance_name
-            else:
+            elif issue:
+                if not re.match(url_regex, issue):
+                    raise ValueError(f"The issue label was not provided or is of incorrect format, example: https://github.com/wazuh/<repository>/issues/<issue-number>")
                 issue_name= re.search(r'github\.com\/wazuh\/([^\/]+)\/issues', issue)
                 repository = cls.generate_repository_name(str(issue_name.group(1)))
                 name = repository + "-" + str(re.search(r'(\d+)$', issue).group(1)) + "-" + str(params.composite_name.split("-")[1]) + "-" + str(params.composite_name.split("-")[2])
+            else:
+                name = str(params.composite_name.split("-")[1]) + "-" + str(params.composite_name.split("-")[2]) + "-" + str(params.composite_name.split("-")[3])
 
             # Keys.
             if platform == "windows":
                 credentials.create_password()
             elif not ssh_key:
                 logger.debug(f"Generating new key pair")
-                credentials.generate(temp_dir, str('-'.join(name.split("-")[:-2])))
+                credentials.generate(temp_dir, name + "-key-" + str(random.randint(1000, 9999)))
             else:
                 logger.debug(f"Using provided key pair")
                 key_id = credentials.ssh_key_interpreter(ssh_key)
@@ -202,7 +206,6 @@ def __create_ec2_instance(config: AWSConfig) -> str:
                 'Tags': [
                     {'Key': 'Name', 'Value': config.name},
                     {'Key': 'termination_date', 'Value': config.termination_date},
-                    {'Key': 'issue', 'Value': config.issue},
                     {'Key': 'team', 'Value': config.team}
                 ]
             }]
@@ -214,6 +217,9 @@ def __create_ec2_instance(config: AWSConfig) -> str:
         if config.host_identifier:
             params['Placement'] = {'AvailabilityZone': config.zone, 'HostId': config.host_identifier}
 
+        if config.issue:
+            params['TagSpecifications'][0]['Tags'].append({'Key': 'issue', 'Value': config.issue})
+
         instance = client.create_instances(**params)[0]
         # Wait until the instance is running.
         instance.wait_until_running()
@@ -239,8 +245,8 @@ def __parse_config(cls, params: CreationPayload, credentials: AWSCredentials, is
         config = {}
 
         # Get the specs from the yamls.
-        size_specs = cls._get_size_specs()[params.size]
-        os_specs = cls._get_os_specs()[params.composite_name]
+        size_specs = cls._get_size_specs(params.size)
+        os_specs = cls._get_os_specs(params.composite_name)
         mics_specs = cls._get_misc_specs()
         arch = params.composite_name.split('-')[-1]
         platform = str(params.composite_name.split("-")[0])

deployability/modules/allocation/generic/provider.py

@@ -154,26 +154,39 @@ def _generate_instance_id(prefix: str) -> str:
         return f"{prefix}-{uuid.uuid4()}".upper()
 
     @classmethod
-    def _get_os_specs(cls) -> dict:
+    def _get_os_specs(cls, composite_name: str) -> dict:
         """
         Gets the OS specifications for the provider.
+        composite_name (str): The name of the composite OS.
 
         Returns:
             dict: A dictionary containing the OS specifications for the provider.
         """
         with open(cls.OS_PATH, "r") as f:
-            return yaml.safe_load(f).get(cls.provider_name)
+            version_available = []
+            os_list = yaml.safe_load(f).get(cls.provider_name)
+            for os in os_list:
+                if os.split("-")[1] == composite_name.split("-")[1]:
+                    version_available.append(os)
+                if str(os) == composite_name:
+                    return os_list[os]
+
+            raise ValueError(f"OS {composite_name} not available for provider {cls.provider_name}. Available versions are {version_available}")
 
     @classmethod
-    def _get_size_specs(cls) -> dict:
+    def _get_size_specs(cls, size: str) -> dict:
         """
         Gets the size specifications for the provider.
+        size (str): The name of the size.
 
         Returns:
             dict: A dictionary containing the size specifications for the provider.
         """
         with open(cls.SIZE_PATH, "r") as f:
-            return yaml.safe_load(f).get(cls.provider_name)
+            size_list = yaml.safe_load(f).get(cls.provider_name)
+            for s in size_list:
+                if size == s:
+                    return size_list[s]
 
     @classmethod
     def _get_misc_specs(cls) -> dict:

deployability/modules/allocation/static/specs/size.yml

@@ -14,7 +14,7 @@ vagrant:
 aws:
   micro:
     amd64:
-      type: t3.micro
+      type: t2.small
     arm64:
       type: a1.medium
   small:

deployability/modules/allocation/vagrant/provider.py

@@ -195,8 +195,8 @@ def __parse_config(cls, params: CreationPayload, credentials: VagrantCredentials
         """
         config = {}
         # Get the specs from the yamls.
-        size_specs = cls._get_size_specs()[params.size]
-        os_specs = cls._get_os_specs()[params.composite_name]
+        size_specs = cls._get_size_specs(params.size)
+        os_specs = cls._get_os_specs(params.composite_name)
         # Parse the configuration.
         config['ip'] = cls.__get_available_ip()
         config['box'] = str(os_specs['box'])
@@ -303,16 +303,14 @@ def __macos_host(arch: str, action: str) -> str:
                 ssh_password = client.get_secret_value(SecretId='devops_black_mini_jenkins_password')['SecretString']
                 ssh_user = client.get_secret_value(SecretId='devops_black_mini_jenkins_user')['SecretString']
             except Exception as e:
-                logger.error('Could not get macOS Black mini server IP: ' + str(e) + '.')
-                exit(1)
+                raise ValueError('Could not get macOS Black mini server IP: ' + str(e) + '.')
 
             try:
                 tn = Telnet(server_ip, server_port, timeout)
                 conn_ok = True
                 tn.close()
             except Exception as e:
-                logger.error('Could not connect to macOS Black mini server: ' + str(e) + '.')
-                exit(1)
+                raise ValueError('Could not connect to macOS Black mini server: ' + str(e) + '.')
 
             macos_host_parameters['server_ip'] = server_ip
             macos_host_parameters['ssh_password'] = ssh_password
@@ -332,7 +330,6 @@ def __macos_host(arch: str, action: str) -> str:
                         logger.info(f"Using the black mini server to deploy.")
                         return macos_host_parameters
                     else:
-                        logger.error(f"Black mini server is under heavy load, use AWS provider.")
-                        exit(1)
+                        raise ValueError(f"Black mini server is under heavy load, use AWS provider.")
                 else:
                     return macos_host_parameters