Vagrant Inventory Plugin

# pylint: disable=super-with-arguments
#!/usr/bin/env python3
import os
import yaml
from ansible.errors import AnsibleError
from ansible.plugins.inventory import BaseInventoryPlugin
import configparser
import re


class InventoryModule(BaseInventoryPlugin):
    NAME = "vagrant"

    def verify_file(self, path):
        return path.endswith((".yml", ".yaml")) and "Vagrantfile" in os.path.basename(
            path
        )

    def parse(self, inventory, loader, path, cache=True):
        super(InventoryModule, self).parse(inventory, loader, path, cache)
        try:
            with open(path, "r", encoding="utf-8") as file_handle:
                data = yaml.safe_load(file_handle)
        except Exception as exception_error:
            raise AnsibleError(
                f"Error loading Vagrantfile: {str(exception_error)}"
            ) from exception_error
        boxes = data.get("boxes")
        defaults = data.get("defaults", {})
        default_prefix = defaults.get("prefix")
        for node in data.get("nodes", []):
            nodename = node.get("name")
            host_prefix = node.get("prefix", default_prefix)
            host = f"{host_prefix}-{nodename}"
            if not host:
                continue
            self.inventory.add_host(host)
            ip_address = node.get("ip-address")
            if ip_address:
                self.inventory.set_variable(host, "ansible_host", ip_address)
            # Determine environment group based on last letter of prefix
            if host_prefix:
                last_char = host_prefix[-1].lower()
                env_group = None
                if last_char == "d":
                    env_group = "development"
                elif last_char == "t":
                    env_group = "test"
                elif last_char == "a":
                    env_group = "acceptance"
                elif last_char == "p":
                    env_group = "production"
                if env_group:
                    self.inventory.add_group(env_group)
                    self.inventory.add_child(env_group, host)
            # Labels
            node_labels = node.get("labels", [])
            box_labels = boxes.get(node.get("box", ""), {}).get("labels", [])
            for label in node_labels + box_labels:
                self.inventory.add_group(label)
                if label != host:
                    self.inventory.add_child(label, host)

        # Handle additional inventory files
        inventory_files = defaults.get("inventory", [])
        for inv_file in inventory_files:
            inv_path = os.path.join(os.path.dirname(path), inv_file)
            if os.path.exists(inv_path):
                self._parse_ini(inv_path)

        # Export to INI if configured
        inventory_export = defaults.get("inventory_export")
        if inventory_export:
            export_path = os.path.join(os.path.dirname(path), inventory_export)
            self._write_ini(export_path)

    def _parse_ini(self, ini_path):
        print(f"ini_path: {ini_path}")
        config = configparser.ConfigParser(allow_no_value=True)
        with open(ini_path, "r", encoding="utf-8") as f:
            config.read_file(f)

        for section in config.sections():
            if ":" in section:
                group_name, section_type = section.split(":", 1)
                group_name = group_name.strip()
                section_type = section_type.strip()
                if section_type == "children":
                    if group_name not in self.inventory.groups:
                        self.inventory.add_group(group_name)
                    for child in config.options(section):
                        child = child.strip()
                        if child:
                            self.inventory.add_group(child)
                            self.inventory.add_child(group_name, child)
                elif section_type == "vars":
                    if group_name not in self.inventory.groups:
                        self.inventory.add_group(group_name)
                    for var, value in config.items(section):
                        value = value.strip() if value else None
                        self.inventory.set_variable(group_name, var, value)
            else:
                # Regular group
                self.inventory.add_group(section)
                for host in config.options(section):
                    host = host.strip()
                    if host:
                        # Check if it's a host with vars, but in simple ini, options are hosts
                        # If there's value, it might be port or something, but for now assume host
                        self.inventory.add_host(host, group=section)
                        # If there are vars like host=var, but in standard ini, hosts are keys with no value

        # Handle host vars if any, but in standard ini, host vars are in [host:vars] which we handle above

    def _write_ini(self, ini_path):
        config = configparser.ConfigParser(allow_no_value=True)

        # Add groups and hosts
        for group_name, group in self.inventory.groups.items():
            if group_name == 'all' or group_name == 'ungrouped':
                continue
            config.add_section(group_name)
            for host in group.hosts:
                config.set(group_name, host.name, None)

        # Add children
        for group_name, group in self.inventory.groups.items():
            if group.child_groups:
                section = f"{group_name}:children"
                config.add_section(section)
                for child in group.child_groups:
                    config.set(section, child.name, None)

        # Add group vars
        for group_name, group in self.inventory.groups.items():
            vars_dict = group.get_vars()
            if vars_dict:
                section = f"{group_name}:vars"
                config.add_section(section)
                for k, v in vars_dict.items():
                    config.set(section, k, str(v))

        # Add host vars
        for host_name, host in self.inventory.hosts.items():
            vars_dict = host.get_vars()
            # Remove ansible_host as it's usually set separately
            vars_dict = {k: v for k, v in vars_dict.items() if k != 'ansible_host'}
            if vars_dict:
                section = f"{host_name}:vars"
                config.add_section(section)
                for k, v in vars_dict.items():
                    config.set(section, k, str(v))

        # Add [host] sections for hosts with ansible_host
        for host_name, host in self.inventory.hosts.items():
            ansible_host = host.get_vars().get('ansible_host')
            if ansible_host:
                section = host_name
                if not config.has_section(section):
                    config.add_section(section)
                config.set(section, f'ansible_host={ansible_host}')

        with open(ini_path, 'w', encoding='utf-8') as f:
            config.write(f)