nixos-servers/scripts/create-host/manipulators.py

"""Text manipulation for flake.nix and Terraform files."""

import re
from pathlib import Path

from models import HostConfig


def update_flake_nix(config: HostConfig, repo_root: Path, force: bool = False) -> None:
    """
    Add or update host entry in flake.nix nixosConfigurations.

    Args:
        config: Host configuration
        repo_root: Path to repository root
        force: If True, replace existing entry; if False, insert new entry
    """
    flake_path = repo_root / "flake.nix"
    content = flake_path.read_text()

    # Create new entry
    new_entry = f"""      {config.hostname} = nixpkgs.lib.nixosSystem {{
        inherit system;
        specialArgs = {{
          inherit inputs self sops-nix;
        }};
        modules = [
          (
            {{ config, pkgs, ... }}:
            {{
              nixpkgs.overlays = commonOverlays;
            }}
          )
          ./hosts/{config.hostname}
          sops-nix.nixosModules.sops
        ];
      }};
"""

    # Check if hostname already exists
    hostname_pattern = rf"^      {re.escape(config.hostname)} = nixpkgs\.lib\.nixosSystem"
    existing_match = re.search(hostname_pattern, content, re.MULTILINE)

    if existing_match and force:
        # Replace existing entry
        # Match the entire block from "hostname = " to "};"
        replace_pattern = rf"^      {re.escape(config.hostname)} = nixpkgs\.lib\.nixosSystem \{{.*?^      \}};\n"
        new_content, count = re.subn(replace_pattern, new_entry, content, flags=re.MULTILINE | re.DOTALL)

        if count == 0:
            raise ValueError(f"Could not find existing entry for {config.hostname} in flake.nix")
    else:
        # Insert new entry before closing brace of nixosConfigurations
        # Pattern: "      };\n      packages = forAllSystems"
        pattern = r"(      \};)\n(      packages = forAllSystems)"
        replacement = rf"{new_entry}\g<1>\n\g<2>"

        new_content, count = re.subn(pattern, replacement, content)

        if count == 0:
            raise ValueError(
                "Could not find insertion point in flake.nix. "
                "Looking for pattern: '      };\\n      packages = forAllSystems'"
            )

    flake_path.write_text(new_content)


def update_terraform_vms(config: HostConfig, repo_root: Path, force: bool = False) -> None:
    """
    Add or update VM entry in terraform/vms.tf locals.vms map.

    Args:
        config: Host configuration
        repo_root: Path to repository root
        force: If True, replace existing entry; if False, insert new entry
    """
    terraform_path = repo_root / "terraform" / "vms.tf"
    content = terraform_path.read_text()

    # Create new entry based on whether we have static IP or DHCP
    if config.is_static_ip:
        new_entry = f'''    "{config.hostname}" = {{
      ip        = "{config.ip}"
      cpu_cores = {config.cpu}
      memory    = {config.memory}
      disk_size = "{config.disk}"
    }}
'''
    else:
        new_entry = f'''    "{config.hostname}" = {{
      cpu_cores = {config.cpu}
      memory    = {config.memory}
      disk_size = "{config.disk}"
    }}
'''

    # Check if hostname already exists
    hostname_pattern = rf'^\s+"{re.escape(config.hostname)}" = \{{'
    existing_match = re.search(hostname_pattern, content, re.MULTILINE)

    if existing_match and force:
        # Replace existing entry
        # Match the entire block from "hostname" = { to }
        replace_pattern = rf'^\s+"{re.escape(config.hostname)}" = \{{.*?^\s+\}}\n'
        new_content, count = re.subn(replace_pattern, new_entry, content, flags=re.MULTILINE | re.DOTALL)

        if count == 0:
            raise ValueError(f"Could not find existing entry for {config.hostname} in terraform/vms.tf")
    else:
        # Insert new entry before closing brace
        # Pattern: "  }\n\n  # Compute VM configurations"
        pattern = r"(  \})\n\n(  # Compute VM configurations)"
        replacement = rf"{new_entry}\g<1>\n\n\g<2>"

        new_content, count = re.subn(pattern, replacement, content)

        if count == 0:
            raise ValueError(
                "Could not find insertion point in terraform/vms.tf. "
                "Looking for pattern: '  }\\n\\n  # Compute VM configurations'"
            )

    terraform_path.write_text(new_content)


def add_wrapped_token_to_vm(hostname: str, wrapped_token: str, repo_root: Path) -> None:
    """
    Add or update the vault_wrapped_token field in an existing VM entry.

    Args:
        hostname: Hostname of the VM
        wrapped_token: The wrapped token to add
        repo_root: Path to repository root
    """
    terraform_path = repo_root / "terraform" / "vms.tf"
    content = terraform_path.read_text()

    # Find the VM entry
    hostname_pattern = rf'^\s+"{re.escape(hostname)}" = \{{'
    match = re.search(hostname_pattern, content, re.MULTILINE)

    if not match:
        raise ValueError(f"Could not find VM entry for {hostname} in terraform/vms.tf")

    # Find the full VM block
    block_pattern = rf'(^\s+"{re.escape(hostname)}" = \{{)(.*?)(^\s+\}})'
    block_match = re.search(block_pattern, content, re.MULTILINE | re.DOTALL)

    if not block_match:
        raise ValueError(f"Could not parse VM block for {hostname}")

    block_start = block_match.group(1)
    block_content = block_match.group(2)
    block_end = block_match.group(3)

    # Check if vault_wrapped_token already exists
    if "vault_wrapped_token" in block_content:
        # Update existing token
        block_content = re.sub(
            r'vault_wrapped_token\s*=\s*"[^"]*"',
            f'vault_wrapped_token = "{wrapped_token}"',
            block_content
        )
    else:
        # Add new token field (add before closing brace)
        # Find the last field and add after it
        block_content = block_content.rstrip()
        if block_content and not block_content.endswith("\n"):
            block_content += "\n"
        block_content += f'      vault_wrapped_token = "{wrapped_token}"\n'

    # Reconstruct the block
    new_block = block_start + block_content + block_end

    # Replace in content
    new_content = re.sub(
        rf'^\s+"{re.escape(hostname)}" = \{{.*?^\s+\}}',
        new_block,
        content,
        flags=re.MULTILINE | re.DOTALL
    )

    terraform_path.write_text(new_content)