Merge pull request #5623 from wazuh/change/5614-change-allocator-logs…

…-to-reference-tasks-instead-of-generic-values

Add task information in the allocation logs when create or delete an instance

deployability/modules/allocation/allocation.py

@@ -13,7 +13,7 @@
 
 from .aws.provider import AWSProvider, AWSConfig
 from .generic import Instance, Provider, models
-from .generic.utils import logger
+from .generic.utils import logger, logger_with_instance_name
 from .vagrant.provider import VagrantProvider, VagrantConfig
 
 
@@ -58,6 +58,9 @@ def __create(cls, payload: models.CreationPayload):
         config = cls.___get_custom_config(payload)
         instance = provider.create_instance(
             payload.working_dir, instance_params, config, payload.ssh_key)
+
+        global logger
+        logger = logger_with_instance_name(instance)
         logger.info(f"Instance {instance.identifier} created.")
         # Start the instance.
         instance.start()
@@ -74,10 +77,10 @@ def __create(cls, payload: models.CreationPayload):
                 logger.warning(f"Rolling back instance creation.")
                 track_payload = {'track_output': track_file}
                 cls.__delete(track_payload)
-                logger.info(f"Instance {instance.identifier} deleted.")
             else:
                 logger.warning(f'The VM will not be automatically removed. Please remove it executing Allocation module with --action delete.')
-        logger.info(f"Instance {instance.identifier} created successfully.")
+        else:
+            logger.info(f"Instance {instance.identifier} created successfully.")
 
     @classmethod
     def __delete(cls, payload: models.InstancePayload) -> None:
@@ -94,6 +97,9 @@ def __delete(cls, payload: models.InstancePayload) -> None:
             track = models.TrackOutput(**yaml.safe_load(f))
         provider = PROVIDERS[track.provider]()
         provider.destroy_instance(models.InstancePayload(**dict(track)))
+
+        global logger
+        logger = logger_with_instance_name(track)
         logger.info(f"Instance {track.identifier} deleted.")
 
     @staticmethod
@@ -176,6 +182,7 @@ def __generate_track_file(instance: Instance, provider_name: str) -> None:
             inventory = models.InventoryOutput(**yaml.safe_load(f))
         port = inventory.ansible_port
         track = models.TrackOutput(identifier=instance.identifier,
+                                    name=instance.name,
                                     provider=provider_name,
                                     instance_dir=str(instance.path),
                                     key_path=str(instance.credentials.key_path),

deployability/modules/allocation/aws/provider.py

@@ -12,7 +12,7 @@
 import subprocess
 
 from modules.allocation.generic import Provider
-from modules.allocation.generic.models import CreationPayload, InstancePayload, InstancePayload
+from modules.allocation.generic.models import CreationPayload, InstancePayload
 from modules.allocation.generic.utils import logger
 from .credentials import AWSCredentials
 from .instance import AWSInstance
@@ -129,6 +129,7 @@ def _create_instance(cls, base_dir: Path, params: CreationPayload, config: AWSCo
 
         instance_params = {}
         instance_params['instance_dir'] = instance_dir
+        instance_params['name'] = config.name
         instance_params['identifier'] = instance_id
         instance_params['platform'] = platform
         instance_params['host_identifier'] = host_identifier

deployability/modules/allocation/generic/instance.py

@@ -4,7 +4,6 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from .utils import logger
 
 from .credentials import Credentials
 from .models import ConnectionInfo, InstancePayload
@@ -43,6 +42,7 @@ def __init__(self, instance_parameters: InstancePayload, credentials: Credential
 
         self.path: Path = path
         self.identifier: str = str(instance_parameters.identifier)
+        self.name: str = instance_parameters.name
         self.credentials: Credentials = credentials
         self.host_identifier: str = instance_parameters.host_identifier
         self.host_instance_dir: Path = instance_parameters.host_instance_dir

deployability/modules/allocation/generic/models.py

@@ -36,6 +36,7 @@ class InventoryOutput(BaseModel):
 
 class TrackOutput(BaseModel):
     identifier: str
+    name: str
     provider: str
     instance_dir: str
     key_path: str
@@ -116,6 +117,7 @@ class TrackPayload(BaseModel):
 
 class InstancePayload(BaseModel):
     identifier: str
+    name: str | None = None
     instance_dir: str | Path
     key_path: Path | None = None
     host_identifier: str  | None = None

deployability/modules/allocation/generic/utils.py

@@ -2,6 +2,27 @@
 # Created by Wazuh, Inc. <[email protected]>.
 # This program is a free software; you can redistribute it and/or modify it under the terms of GPLv2
 
+import logging
+from modules.allocation.generic.instance import Instance
+from modules.allocation.generic.models import TrackOutput
 from modules.generic.logger import Logger
 
+# Default logger
 logger = Logger("allocator").get_logger()
+
+
+def logger_with_instance_name(instance_info: Instance | TrackOutput) -> logging.Logger:
+    """
+    Returns a logger with the instance name if it is different from the identifier,
+    otherwise returns the default logger without the name.
+
+    Args:
+        instance (Instance): The instance object.
+
+    Returns:
+        logging.Logger: The logger object.
+    """
+
+    if instance_info.name != instance_info.identifier:
+        return Logger(f"allocator [{instance_info.name}]").get_logger()
+    return logger

deployability/modules/allocation/vagrant/provider.py

@@ -121,6 +121,7 @@ def _create_instance(cls, base_dir: Path, params: CreationPayload, config: Vagra
         instance_params = {}
         instance_params['instance_dir'] = instance_dir
         instance_params['identifier'] = instance_id
+        instance_params['name'] = config.name
         instance_params['platform'] = platform
         instance_params['host_identifier'] = host_identifier
         instance_params['host_instance_dir'] = host_instance_dir