Improve on allocation

deployability/launchers/allocation.py

@@ -15,6 +15,7 @@ def parse_arguments():
     parser.add_argument("--composite-name", required=False, default=None)
     parser.add_argument("--action", choices=['create', 'delete'], required=False, default='create')
     parser.add_argument("--custom-credentials", required=False, default=None)
+    parser.add_argument("--custom-provider-config", required=False, default=None)
     parser.add_argument("--track-output", required=False, default='/tmp/wazuh-qa/track.yml')
     parser.add_argument("--inventory-output", required=False, default='/tmp/wazuh-qa/inventory.yml')
     parser.add_argument("--working-dir", required=False, default='/tmp/wazuh-qa')

deployability/modules/__init__.py

@@ -1,5 +1,4 @@
 from .provision import Provision
 from .generic import Ansible
-from .allocation.vagrant import VagrantProvider
-from .allocation.aws import AWSProvider
+from .allocation import Allocator
 from .generic import SchemaValidator

deployability/modules/allocation/allocation.py

@@ -3,10 +3,12 @@
 from pathlib import Path
 
 from .generic import Instance, Provider, models
-from .aws.provider import AWSProvider
-from .vagrant.provider import VagrantProvider
+from .aws.provider import AWSProvider, AWSConfig
+from .vagrant.provider import VagrantProvider, VagrantConfig
+
 
 PROVIDERS = {'vagrant': VagrantProvider, 'aws': AWSProvider}
+CONFIGS = {'vagrant': VagrantConfig, 'aws': AWSConfig}
 
 
 class Allocator:
@@ -39,12 +41,13 @@ def __create(cls, payload: models.CreationPayload):
 
         Args:
             payload (CreationPayload): The payload containing the parameters
-                                              for instance creation.
+                                       for instance creation.
         """
         instance_params = models.CreationPayload(**dict(payload))
         provider: Provider = PROVIDERS[payload.provider]()
+        config = cls.___get_custom_config(payload)
         instance = provider.create_instance(
-            payload.working_dir, instance_params)
+            payload.working_dir, instance_params, config)
         print(f"Instance {instance.identifier} created.")
         # Start the instance
         instance.start()
@@ -70,6 +73,26 @@ def __delete(cls, payload: models.DeletionPayload) -> None:
         provider.destroy_instance(track.instance_dir, track.identifier)
         print(f"Instance {track.identifier} deleted.")
 
+    @staticmethod
+    def ___get_custom_config(payload: models.CreationPayload) -> models.ProviderConfig | None:
+        """
+        Gets the custom configuration from a file.
+
+        Args:
+            payload (CreationPayload): The payload containing the parameters
+                                       for instance creation.
+
+        Returns:
+            ProviderConfig: The configuration object.
+        """
+        if not payload.custom_provider_config:
+            return None
+        # Read the custom config file and validate it.
+        config_model: models.ProviderConfig = CONFIGS[payload.provider]
+        with open(payload.custom_provider_config, 'r') as f:
+            config = config_model(**yaml.safe_load(f))
+        return config
+
     @staticmethod
     def __generate_inventory(instance: Instance, inventory_path: Path) -> None:
         """

deployability/modules/allocation/aws/provider.py

@@ -20,40 +20,43 @@ class AWSProvider(Provider):
     """
 
     provider_name = 'aws'
-    _client = boto3.resource('ec2')
 
     @classmethod
-    def _create_instance(cls, base_dir: Path, params: CreationPayload) -> AWSInstance:
+    def _create_instance(cls, base_dir: Path, params: CreationPayload, config: AWSConfig = None) -> AWSInstance:
         """
         Create an AWS EC2 instance.
 
         Args:
             base_dir (Path): Base directory for storing instance data.
             params (CreationPayload): Payload containing creation parameters.
+            config (AWSConfig, optional): Configuration for the instance. Defaults to None.
 
         Returns:
             AWSInstance: Created AWSInstance object.
         """
         temp_id = cls._generate_instance_id(cls.provider_name)
         temp_dir = base_dir / temp_id
-        # Generate the credentials.
         credentials = AWSCredentials()
-        credentials.generate(temp_dir, temp_id.split('-')[-1] + '_key')
-        # Parse the config and create the AWS EC2 instance.
-        config = cls._parse_config(params, credentials)
-        _instance = cls._client.create_instances(ImageId=config.ami,
-                                                 InstanceType=config.type,
-                                                 KeyName=config.key_name,
-                                                 SecurityGroupIds=config.security_groups,
-                                                 MinCount=1, MaxCount=1)[0]
-        # Wait until the instance is running.
-        _instance.wait_until_running()  
+        if not config:
+            # Generate the credentials.
+            credentials.generate(temp_dir, temp_id.split('-')[-1] + '_key')
+            # Parse the config if it is not provided.
+            config = cls.__parse_config(params, credentials)
+        else:
+            # Load the existing credentials.
+            credentials.load(config.key_name)
+            # Create the temp directory. 
+            # TODO: Review this on the credentials refactor.
+            if not temp_dir.exists():
+                temp_dir.mkdir(parents=True, exist_ok=True)
+        # Generate the instance.
+        instance_id = cls.__create_ec2_instance(config)
         # Rename the temp directory to its real name.
-        instance_dir = Path(base_dir, _instance.instance_id)
+        instance_dir = Path(base_dir, instance_id)
         os.rename(temp_dir, instance_dir)
         credentials.key_path = (instance_dir / credentials.name).with_suffix('.pem')
 
-        return AWSInstance(instance_dir, _instance.instance_id, credentials, config.user)
+        return AWSInstance(instance_dir, instance_id, credentials, config.user)
 
     @staticmethod
     def _load_instance(instance_dir: Path, instance_id: str) -> AWSInstance:
@@ -83,8 +86,29 @@ def _destroy_instance(cls, instance_dir: str, identifier: str) -> None:
             instance.credentials.delete()
         instance.delete()
 
+    @staticmethod
+    def __create_ec2_instance(config: AWSConfig) -> str:
+        """
+        Create an AWS EC2 instance.
+
+        Args:
+            config (AWSConfig): Configuration for the instance.
+
+        Returns:
+            str: Identifier of the created instance.
+        """
+        client = boto3.resource('ec2')
+        instance = client.create_instances(ImageId=config.ami,
+                                           InstanceType=config.type,
+                                           KeyName=config.key_name,
+                                           SecurityGroupIds=config.security_groups,
+                                           MinCount=1, MaxCount=1)[0]
+        # Wait until the instance is running.
+        instance.wait_until_running()
+        return instance.instance_id
+
     @classmethod
-    def _parse_config(cls, params: CreationPayload, credentials: AWSCredentials) -> AWSConfig:
+    def __parse_config(cls, params: CreationPayload, credentials: AWSCredentials) -> AWSConfig:
         """
         Parse configuration parameters for creating an AWS EC2 instance.
 

deployability/modules/allocation/generic/models.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from pydantic import BaseModel, IPvAnyAddress, field_validator
+from pydantic import BaseModel, IPvAnyAddress, field_validator, model_validator
 from typing_extensions import Literal
 
 
@@ -35,32 +35,49 @@ class TrackOutput(BaseModel):
 
 class InputPayload(BaseModel):
     action: Literal['create', 'delete', 'status'] = 'create'
-    provider: str | None
-    size: Literal['micro', 'small', 'medium', 'large', None]
-    composite_name: str | None
-    track_output: Path | None
-    inventory_output: Path | None
-    working_dir: Path | None
-    custom_credentials: str | None
+    provider: str | None = None
+    size: Literal['micro', 'small', 'medium', 'large', None] = None
+    composite_name: str | None = None
+    working_dir: Path | None = Path('/tmp/wazuh-qa')
+    track_output: Path | None = working_dir / 'track.yml'
+    inventory_output: Path | None = working_dir / 'inventory.yml'
+    custom_credentials: str | None = None
+    custom_provider_config: Path | None = None
 
 
 class CreationPayload(InputPayload):
     provider: str
-    size: Literal['micro', 'small', 'medium', 'large']
-    composite_name: str
+    size: Literal['micro', 'small', 'medium', 'large'] | None = None
+    composite_name: str | None = None
     track_output: Path
     inventory_output: Path
     working_dir: Path
     custom_credentials: str | None = None
-
-    @field_validator('custom_credentials')
+    custom_provider_config: Path | None = None
+
+    @model_validator(mode='before')
+    def validate_dependency(cls, values) -> dict:
+        """Validate required fields."""
+        required_if_not_config = ['composite_name', 'size']
+        if values.get('custom_provider_config'):
+            return values
+        for attr in required_if_not_config:
+            if not values.get(attr):
+                raise ValueError(f"{attr} is required if custom_provider_config is not provided.")
+
+        return values
+
+    @field_validator('custom_provider_config')
     @classmethod
-    def check_credentials(cls, v: str) -> str | None:
+    def check_config(cls, v: Path | None) -> Path | None:
         if not v:
             return None
-        path = Path(v)
-        if not path.exists() or not path.is_file():
-            raise ValueError(f"Invalid credentials path: {path}")
+        if not v.exists():
+            raise ValueError(f"Custom provider config file does not exist: {v}")
+        elif not v.is_file():
+            raise ValueError(f"Custom provider config file is not a file: {v}")
+        elif not v.suffix in ['.yml', '.yaml']:
+            raise ValueError(f"Custom provider config file must be yaml: {v}")
         return v
 
     @field_validator('working_dir', mode='before')

deployability/modules/allocation/generic/provider.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 
 from .instance import Instance
-from .models import CreationPayload
+from .models import CreationPayload, ProviderConfig
 
 
 class Provider(ABC):
@@ -47,20 +47,21 @@ def provider_name(self) -> str:
         pass
 
     @classmethod
-    def create_instance(cls, base_dir: str | Path, params: CreationPayload) -> Instance:
+    def create_instance(cls, base_dir: str | Path, params: CreationPayload, config: ProviderConfig = None) -> Instance:
         """
         Creates a new instance.
 
         Args:
             base_dir (str | Path): The base directory for the instance.
             params (CreationPayload): The parameters for creating the instance.
+            config (ProviderConfig, optional): The configuration for the instance. Defaults to None.
 
         Returns:
             Instance: The created instance.
         """
         params = CreationPayload(**dict(params))
         base_dir = Path(base_dir)
-        return cls._create_instance(base_dir, params)
+        return cls._create_instance(base_dir, params, config)
 
     @classmethod
     def load_instance(cls, instance_dir: str | Path, instance_id: str) -> Instance:
@@ -97,13 +98,14 @@ def destroy_instance(cls, instance_dir: str | Path, identifier: str) -> None:
 
     @classmethod
     @abstractmethod
-    def _create_instance(cls, base_dir: Path, params: CreationPayload) -> Instance:
+    def _create_instance(cls, base_dir: Path, params: CreationPayload, config: ProviderConfig = None) -> Instance:
         """
         Abstract method that creates a new instance.
 
         Args:
             base_dir (Path): The base directory for the instance.
             params (CreationPayload): The parameters for creating the instance.
+            config (ProviderConfig, optional): The configuration for the instance. Defaults to None.
 
         Returns:
             Instance: The created instance.

deployability/modules/allocation/static/templates/vagrant.j2

@@ -1,10 +1,30 @@
 Vagrant.configure("2") do |config|
+    # Box image settings
     config.vm.box = "{{ config.box }}"
     config.vm.box_version = "{{ config.box_version }}"
+
+    # Copy the public key to the VM
     config.vm.provision "file", source: "{{ config.public_key }}", destination: ".ssh/authorized_keys"
-    config.vm.network "private_network", ip:"{{ config.ip }}"
+
+    # VirtualBox specific settings
     config.vm.provider "virtualbox" do |v|
         v.memory = {{ config.memory }}
         v.cpus = {{ config.cpu }}
     end
+
+    # Network settings
+    config.vm.network "private_network", ip:"{{ config.ip }}"
+    config.ssh.forward_agent = true
+    # Create a file to indicate that the VM has been initialized
+    # This is required to correctly get the ssh-config of the VM
+    # TODO: find a better solution
+    init_indicator = "./init"
+    if not ::File.exists?(init_indicator)
+        File.write(init_indicator, "initialized")
+    else
+        # Use the IP address of the VM to connect via SSH
+        config.ssh.host = "{{ config.ip }}"
+        config.ssh.port = 22
+        config.ssh.private_key_path = "{{ config.private_key }}"
+    end
 end

deployability/modules/allocation/vagrant/credentials.py

@@ -71,30 +71,22 @@ def generate(self, base_dir: str | Path, name: str) -> Path:
         self.key_path = private_key_path
         return self.key_path
 
-    def load(self, base_dir: str | Path, name: str) -> None:
+    def load(self, path: str | Path) -> None:
         """
         Loads an existing key pair from the specified directory.
 
         Args:
-            base_dir (str | Path): The directory where the key pair is stored.
-            name (str): The filename of the key pair.
+            path (str | Path): The path to the key pair.
 
         Raises:
             CredentialsError: This exception is raised if the key pair doesn't exist or the specified directory is invalid.
         """
-        base_dir = Path(base_dir)
-        if not base_dir.exists() or not base_dir.is_dir():
-            raise self.CredentialsError(f"Invalid path {base_dir}.")
-        key_path = Path(base_dir, name)
-        pub_key_path = key_path.with_suffix(".pub")
+        key_path = Path(path)
         if not key_path.exists() or not key_path.is_file():
-            raise self.CredentialsError(f"Invalid key name {name}.")
-        if not pub_key_path.exists() or not pub_key_path.is_file():
-            raise self.CredentialsError(f"Non-existen public key for {name}.")
-        # Save instance attributes.
+            raise self.CredentialsError(f"Invalid path {key_path}.")
         self.key_path = key_path
-        self.name = name
-        self.key_id = name
+        self.name = key_path.name
+        self.key_id = key_path.name
 
     def delete(self) -> None:
         """

deployability/modules/allocation/vagrant/models.py

@@ -10,6 +10,7 @@ class VagrantConfig(ProviderConfig):
     memory: int
     box: str
     box_version: str
+    private_key: str
     public_key: str
 
     @field_validator('public_key', mode='before')