DTT1 - Allocator module - Code refactor

[QU3B1M]

EPIC: #4495

Description

After finishing the module for the POC, it was determined that the code of this module could be more efficiently diagramed.

This issue aims to analyze & improve the current code of the allocator module, and also define a standard structure for all the modules directories.

New dir structure

/allocator
├── allocation.py
├── aws
│   ├── credentials.py
│   ├── __init__.py
│   ├── instance.py
│   ├── models.py
│   └── provider.py
├── generic
│   ├── credentials.py
│   ├── __init__.py
│   ├── instance.py
│   ├── models.py
│   └── provider.py
├── __init__.py
├── static
│   ├── specs
│   │   ├── misc.yml
│   │   ├── os.yml
│   │   └── size.yml
│   └── templates
│       └── vagrant.j2
└── vagrant
    ├── credentials.py
    ├── __init__.py
    ├── instance.py
    ├── models.py
    └── provider.py

Usage

1. Install dependencies

   pip install -r deployability/de[s/requirements.txt

2. Launch allocation tool
   • Create instance

     python deployability/launchers/allocation.py --action create --provider aws --size small --composite-name linux-ubuntu-22.04-amd64 [OPTIONALS: --track-output, --inventory-output, --working-dir]

     It will output an inventory.yaml with the instance's connection details, and a track_output.yaml file with the instance tracking details (this is used by the deletion action)

   • Delete instance

     python deployability/launchers/allocation.py --action delete --track-output

[QU3B1M]

Update report

• Generated a basic classes diagram with the new approach
  

[QU3B1M]

Update report

• Added pydantic models to validate Inventory & Instance data structures
• Created ProviderInterface to ensure the basic behavior of all the provider classes
• Working on VagrantProvider to handle the vagrant VMs allocation

[QU3B1M]

Update report

• Credential generic class and VagrantCredential specific class created and working correctly

  class Credential(ABC):
      """Interface for Credentials"""
      class KeyCreationError(Exception):
          pass
  
      def __init__(self, base_dir: str | Path, name: str):
          """Initialize Credentials"""
          self.base_dir = Path(base_dir)
          self.name = str(name)
  
          self.private_key: Path = None
          self.public_key: Path = None
  
      @abstractmethod
      def generate(self, **kwargs) -> tuple[str, str] | None:
          """Get credentials"""
          raise NotImplementedError()
  
      @abstractmethod
      def delete(self, **kwargs):
          """Set credentials"""
          raise NotImplementedError()
  
  
  class VagrantCredential(Credential):
  
      def __init__(self, base_dir: str | Path, name: str) -> None:
          super().__init__(base_dir, name)
  
      def generate(self, overwrite: bool = False) -> tuple[str, str] | None:
          if self.private_key and self.public_key and not overwrite:
              return str(self.private_key), str(self.public_key)
          if not self.base_dir.exists():
              self.base_dir.mkdir(parents=True, exist_ok=True)
  
          path = self.base_dir / self.name
          command = ["ssh-keygen",
                     "-f", str(path),
                     "-m", "PEM",
                     "-t", "rsa",
                     "-N", "",
                     "-q"]
          output = subprocess.run(command, check=True, capture_output=True, text=True)
          os.chmod(path, 0o600)
          if output.returncode != 0:
              raise self.KeyCreationError(f"Error creating key pair: {output.stderr}")
  
          self.private_key = path
          self.public_key = path.with_suffix(".pub")
  
          return str(self.private_key), str(self.public_key)
  
      def delete(self) -> None:
          if self.private_key:
              self.private_key.unlink(missing_ok=True)
              self.private_key = None
          if self.public_key:
              self.public_key.unlink(missing_ok=True)
              self.public_key = None

• Continue working on VagrantProvider class

[QU3B1M]

Update report

• Improved credentials classes for a better handling of the key pair

  import os
  import subprocess
  
  from pathlib import Path
  
  from .base import Credentials
  from ..models import CredentialsKeyPair
  
  
  class VagrantCredentials(Credentials):
      """
      A class for generating and deleting Vagrant credentials.
  
      Attributes:
          path (Path): The path to store the credentials.
          name (str): The name of the credentials.
          key_pair (CredentialsKeyPair): The key pair.
  
      Raises:
          KeyCreationError: An error occurred while creating the key.
  
      """
  
      def __init__(self, path: str | Path, name: str) -> None:
          """
          Initializes the VagrantCredentials object.
  
          Args:
              path (Path): The path to store the credentials.
              name (str): The name of the credentials.
          """
          super().__init__(path, name)
  
      def generate_key_pair(self, overwrite: bool = False) -> CredentialsKeyPair:
          """
          Generates a new key pair and returns it.
  
          Args:
              overwrite (bool): Whether to overwrite the existing key pair. Defaults to False.
  
          Returns:
              CredentialsKeyPair: The paths of the key pair.
  
          Raises:
              KeyCreationError: An error occurred while creating the key.
  
          """
          if self.key_pair and not overwrite:
              return self.key_pair
  
          private_key_path = Path(self.path / self.name)
          public_key_path = private_key_path.with_suffix(".pub")
          # Create the base directory if it doesn't exist.
          if not self.path.exists():
              self.path.mkdir(parents=True, exist_ok=True)
          # Delete the existing key pair if it exists.
          if private_key_path.exists():
              private_key_path.unlink()
          if public_key_path.exists():
              public_key_path.unlink()
  
          command = ["ssh-keygen",
                     "-f", str(private_key_path),
                     "-m", "PEM",
                     "-t", "rsa",
                     "-N", "",
                     "-q"]
          output = subprocess.run(command, check=True,
                                  capture_output=True, text=True)
          os.chmod(private_key_path, 0o600)
          if output.returncode != 0:
              raise self.KeyCreationError(
                  f"Error creating key pair: {output.stderr}")
          self.key_pair = CredentialsKeyPair(public_key=str(public_key_path),
                                             private_key=str(private_key_path))
          return self.key_pair
  
      def delete_key_pair(self) -> None:
          """Deletes the key pair."""
          if not self.key_pair:
              return
          Path(self.key_pair.private_key).unlink()
          Path(self.key_pair.public_key).unlink()
          self.key_pair = None

• The instance creation is working correctly on the VagrantProvider class

  import re
  import subprocess
  import uuid
  
  from fnmatch import fnmatch
  from pathlib import Path
  from jinja2 import Environment, FileSystemLoader
  
  from .base import Provider, TEMPLATES_DIR
  from ..credentials.vagrant import VagrantCredentials
  from ..models import CredentialsKeyPair, Instance, InstanceParams, Inventory, VagrantConfig
  
  
  class VagrantProvider(Provider):
      """A class for managing Vagrant providers.
  
      Attributes:
          name (str): The name of the provider.
          provider_name (str): The name of the provider.
          working_dir (Path): The working directory for the provider.
          instance_params (InstanceParams): The instance parameters.
          credentials (CredentialsKeyPair): The credentials key pair paths.
          config (ProviderConfig): The provider configuration.
          instance (Instance): The instance.
          inventory (Inventory): The inventory.
      """
      provider_name = 'vagrant'
  
      def __init__(self, base_dir: Path | str, instance_params: InstanceParams, credentials: CredentialsKeyPair) -> None:
          """Initializes the VagrantProvider object.
  
          Args:
              base_dir (Path): The base directory for the provider.
              credentials (CredentialsKeyPair): The credentials key pair paths.
              instance_params (InstanceParams): The instance parameters.
          """
          super().__init__(base_dir, instance_params, credentials)
  
      def create(self) -> Instance:
          """Creates a new vagrant VM instance.
  
          Returns:
              Instance: The instance specifications.
  
          """
          if self.instance:
              return self.instance
          if not self.working_dir.exists():
              self.working_dir.mkdir(parents=True, exist_ok=True)
  
          # Get the config and the instance definitions.
          self.config = self.__parse_config()
          self.instance = self.__generate_instance()
          # Render and write Vagrantfile
          vagrantfile = self.__render_vagrantfile()
          self.__save_vagrantfile(vagrantfile)
  
          return self.instance
  
      def start(self) -> Inventory:
          """Starts the vagrant VM.
  
          Returns:
              Inventory: The ansible inventory of the instance.
          """
          if not self.instance:
              return
          self.__run_vagrant_command('up')
          self.inventory = self.__parse_inventory()
  
          return self.inventory
  
      def stop(self) -> None:
          """Stops the vagrant VM."""
          if not self.instance:
              return
          self.__run_vagrant_command('halt')
  
      def delete(self) -> None:
          """Deletes the vagrant VM and cleans the environment."""
          if not self.instance:
              return
          self.__run_vagrant_command('destroy -f')
          self.working_dir.rmdir()
          self.inventory = None
          self.instance = None
  
      def status(self) -> str:
          """Checks the status of the vagrant VM.
  
          Returns:
              str: The status of the instance.
          """
          if not self.instance:
              return
          self.__run_vagrant_command('status')
  
      # Internal methods
  
      def __parse_inventory(self) -> Inventory:
          """Parses the inventory info and returns it.
  
          Returns:
              Inventory: The ansible inventory of the instance.
          """
          inventory = {}
          private_key = self.credentials.private_key
          ssh_config = self.__run_vagrant_command('ssh-config')
          patterns = {'ansible_hostname': r'HostName (.*)',
                      'ansible_user': r'User (.*)',
                      'ansible_port': r'Port (.*)'}
          # Parse the inventory.
          inventory['ansible_ssh_private_key_file'] = private_key
          for key, pattern in patterns.items():
              match = re.search(pattern, ssh_config)
              if match:
                  inventory[key] = match.group(1)
              else:
                  raise ValueError(f"Couldn't find {key} in vagrant ssh-config")
  
          return Inventory(**inventory)
  
      def __parse_config(self) -> VagrantConfig:
          """Parses the config and returns it.
  
          Returns:
              VagrantConfig: The vagrant VM configuration.
          """
          config = {}
          composite_name = self.instance_params.composite_name
          roles = self._get_role_specs()
          os_specs = self._get_os_specs()
          # Parse the configuration.
          config['id'] = str(self.__generate_instance_id())
          config['box'] = os_specs[composite_name]['box']
          config['box_version'] = os_specs[composite_name]['box_version']
          for pattern, specs in roles[self.instance_params.role].items():
              if fnmatch(composite_name, pattern):
                  config['cpu'] = specs['cpu']
                  config['memory'] = specs['memory']
                  config['ip'] = specs['ip']
                  break
          return VagrantConfig(**config)
  
      def __generate_instance(self) -> Instance:
          """Generates a new instance.
  
          Returns:
              Instance: The instance specifications.
  
          """
          instance = Instance(name=self.instance_params.name,
                              params=self.instance_params,
                              path=str(self.working_dir),
                              provider=self.provider_name,
                              credential=str(self.credentials.private_key),
                              connection_info=None,
                              provider_config=self.config)
          return instance
  
      def __generate_instance_id(self, prefix: str = "VAGRANT") -> str:
          """
          Generates a random instance id with the given prefix.
  
          Args:
              prefix (str): The prefix for the instance id. Defaults to "VAGRANT".
  
          Returns:
              str: The instance id.
  
          """
          return f"{prefix}-{uuid.uuid4()}".upper()
  
      def __run_vagrant_command(self, command: str) -> str:
          """
          Runs a Vagrant command and returns its output.
  
          Args:
              command (str): The vagrant command to run.
  
          Returns:
              str: The output of the command.
          """
          output = subprocess.run(["vagrant", command],
                                  cwd=self.base_dir,
                                  check=True,
                                  stdout=subprocess.PIPE,
                                  stderr=subprocess.PIPE)
          return output.stdout.decode("utf-8")
  
      def __render_vagrantfile(self) -> str:
          """
          Renders the Vagrantfile template and returns it.
  
          Returns:
              str: The rendered Vagrantfile.
  
          """
          public_key = self.credentials.public_key
          environment = Environment(loader=FileSystemLoader(TEMPLATES_DIR))
          template = environment.get_template("vagrant.j2")
  
          return template.render(config=self.config, credential=public_key)
  
      def __save_vagrantfile(self, vagrantfile: str) -> None:
          """
          Saves the Vagrantfile to disk.
  
          Args:
              vagrantfile (str): The Vagrantfile to save.
  
          """
          with open(self.working_dir / 'Vagrantfile', 'w') as f:
              f.write(vagrantfile)

• Working on a better implementation of the provider interface

[QU3B1M]

Update report

• Improved provider classes
• VagrantProvider class working correctly

  Had some issues running vagrant inside the VM I use for development

• Start the refactor of the AWSProvider class using the generic Provider as base class

[QU3B1M]

Update report

• Create Handler generic and VagrantHandler specific classes, these objects will interact directly with the service (in this case Vagrant or AWS), so the Provider class has less responsabilities

  from pathlib import Path
  import re
  import subprocess
  
  from .generic import ConnectionInfo, Handler
  
  class VagrantHandler(Handler):
      def __init__(self, working_dir: str | Path, vagrantfile_content: str = None) -> None:
          super().__init__(working_dir)
          self.vagrantfile_path = Path(self.working_dir, 'Vagrantfile')
          self.vagrantfile_defined = False
  
          if vagrantfile_content:
              # Write the Vagrantfile with the given content.
              self.write_vagrantfile(vagrantfile_content)
              self.vagrantfile_defined = True
          elif self.read_vagrantfile():
              # It will use the already existing Vagrantifile.
              self.vagrantfile_defined = True
  
      def start(self) -> None:
          """Starts the vagrant VM."""
          if not self.vagrantfile_defined:
              raise Exception('Vagrantfile not defined.')
          self.__run_vagrant_command('up')
  
      def stop(self) -> None:
          """Stops the vagrant VM."""
          if not self.vagrantfile_defined:
              raise Exception('Vagrantfile not defined.')
          self.__run_vagrant_command('halt')
  
      def delete(self) -> None:
          """Deletes the vagrant VM and cleans the environment."""
          if not self.vagrantfile_defined:
              raise Exception('Vagrantfile not defined.')
          self.__run_vagrant_command('destroy -f')
  
      def status(self) -> str:
          """Checks the status of the vagrant VM.
  
          Returns:
              str: The status of the instance.
          """
          if not self.vagrantfile_defined:
              raise Exception('Vagrantfile not defined.')
          output = self.__run_vagrant_command('status')
          return self.__parse_vagrant_status(output)
  
      def get_ssh_config(self) -> ConnectionInfo:
          """Returns the ssh config of the vagrant VM.
  
          Returns:
              ConnectionInfo: The VM's ssh config.
          """
          if not self.vagrantfile_defined:
              raise Exception('Vagrantfile not defined.')
          ssh_config = {}
          output = self.__run_vagrant_command('ssh-config')
          patterns = {'hostname': r'HostName (.*)',
                      'user': r'User (.*)',
                      'port': r'Port (.*)',
                      'private_key': r'IdentityFile (.*)'}
          # Parse the ssh-config.
          for key, pattern in patterns.items():
              match = re.search(pattern, output)
              if match:
                  ssh_config[key] = match.group(1)
              else:
                  raise ValueError(f"Couldn't find {key} in vagrant ssh-config")
          return ConnectionInfo(**ssh_config)
  
      def __run_vagrant_command(self, command: str) -> str:
          """
          Runs a Vagrant command and returns its output.
  
          Args:
              command (str): The vagrant command to run.
  
          Returns:
              str: The output of the command.
          """
          try:
              output = subprocess.run(["vagrant", command],
                                      cwd=self.working_dir,
                                      check=True,
                                      stdout=subprocess.PIPE,
                                      stderr=subprocess.PIPE)
  
              if stderr := output.stderr.decode("utf-8"):
                  print(stderr)
                  print(output.stdout.decode("utf-8"))
              # logging.warning(f"Command '{command}' completed with errors:\n{stderr}")
  
              return output.stdout.decode("utf-8")
  
          except subprocess.CalledProcessError as e:
              print(e)
              # logging.error(f"Command '{command}' failed with error {e.returncode}:\n{e.output.decode('utf-8')}")
              return None
  
      def write_vagrantfile(self, data: str) -> None:
          """
          Saves a Vagrantfile in the current working_dir.
  
          Args:
              data (str): The Vagrantfile content to save.
  
          """
          with open(self.vagrantfile_pat, 'w') as f:
              f.write(data)
          self.vagrantfile_defined = True
  
      def read_vagrantfile(self) -> str:
          """
          Reads the Vagrantfile in the current working_dir.
  
          Returns:
              str: The Vagrantfile content.
  
          """
          if not self.vagrantfile_path.exists():
              return None
          with open(self.vagrantfile_path, 'r') as f:
              return f.read()
  
      def __parse_vagrant_status(self, message: str) -> str:
          lines = message.split('\n')
          for line in lines:
              if 'Current machine states:' in line:
                  status_line = lines[lines.index(line) + 2]
                  status = status_line.split()[1]
                  return status

• Add AmazonEC2Credentials class to handle the generation/deletion of the amazon credentials

  import os
  import boto3
  
  from pathlib import Path
  
  from .generic import Credentials, CredentialsKeyPair
  
  
  class AmazonEC2KeyPair(CredentialsKeyPair):
      id: str
      public_key: str = None
  
  
  class AmazonEC2Credentials(Credentials):
      """
      A class for generating and deleting EC2 credentials.
  
      Attributes:
          path (Path): The path to store the credentials.
          name (str): The name of the credentials.
          key_pair (CredentialsKeyPair): The key pair.
  
      Raises:
          KeyCreationError: An error occurred while creating the key.
  
      """
  
      def __init__(self, path: str | Path, name: str) -> None:
          """
          Initializes the AmazonEC2Credentials object.
  
          Args:
              path (Path): The path to store the credentials.
              name (str): The name of the credentials.
          """
          self._client = boto3.resource('ec2')
          super().__init__(path, name)
  
      def generate_key_pair(self, overwrite: bool = False) -> AmazonEC2KeyPair:
          """
          Generates a new key pair and returns it.
  
          Args:
              overwrite (bool): Whether to overwrite the existing key pair. Defaults to False.
  
          Returns:
              AmazonEC2KeyPair: The paths of the key pair.
  
          """
          if self.key_pair and not overwrite:
              return self.key_pair
          private_key_path = Path(self.path / self.name)
          # Request the key pair from AWS.
          response = self._client.create_key_pair(KeyName=self.name)
          key_pair_id = response.key_pair_id
          with open(private_key_path, 'w') as key_file:
              key_file.write(response.key_material)
          os.chmod(private_key_path, 0o600)
  
          self.key_pair = AmazonEC2KeyPair(private_key=str(private_key_path),
                                           id=key_pair_id)
          return self.key_pair
  
      def delete_key_pair(self) -> None:
          """Deletes the key pair."""
          if not self.key_pair:
              return
          self._client.KeyPair(self.name).delete()
          Path(self.key_pair.private_key).unlink()
          self.key_pair = None

[QU3B1M]

Update report

• Working on the development of the Amazon EC2 provider integration

[QU3B1M]

Update report

• Allocator working with some log prints

Creating instance at C:\tmp\wazuh-qa
Instance VAGRANT-9010BB78-1C17-4C81-BEA8-A7717E783496 created.

Inventory file generated at \tmp\wazuh-qa\inventory.yml

Track file generated at \tmp\wazuh-qa\track.yml

[QU3B1M]

Update report

• The allocator creates and destroys AWS and Vagrant VMs correctly
• Re-ordered the code and improve docstrings

├── allocation.py
├── aws
│   ├── credentials.py
│   ├── __init__.py
│   ├── instance.py
│   ├── models.py
│   └── provider.py
├── generic
│   ├── credentials.py
│   ├── __init__.py
│   ├── instance.py
│   ├── models.py
│   └── provider.py
├── __init__.py
├── static
│   ├── specs
│   │   ├── misc.yml
│   │   ├── os.yml
│   │   └── size.yml
│   └── templates
│       └── vagrant.j2
└── vagrant
    ├── credentials.py
    ├── __init__.py
    ├── instance.py
    ├── models.py
    └── provider.py

[fcaffieri]

LGTM

[fcaffieri]

The unit test will be executed in the PoC Iteration 2#4751