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