azure-vm-sdk 0.6.0__py3-none-any.whl

azure_vm/__init__.py

@@ -0,0 +1,32 @@
+"""Public API for creating and managing Azure VMs through OpenTofu."""
+
+from ._backend import CommandResult
+from .client import AzureClient
+from .exceptions import (
+    AzureVmCommandError,
+    AzureVmError,
+    AzureVmTimeoutError,
+    SshConnectionError,
+    TofuNotInstalledError,
+    VmNotFoundError,
+)
+from .models import ImageInfo, VmArchitecture, VmConfig, VmInfo, VmSize, VmState
+from .vm import AzureVM
+
+__all__ = [
+    "AzureClient",
+    "AzureVM",
+    "AzureVmCommandError",
+    "AzureVmError",
+    "AzureVmTimeoutError",
+    "CommandResult",
+    "ImageInfo",
+    "SshConnectionError",
+    "TofuNotInstalledError",
+    "VmArchitecture",
+    "VmConfig",
+    "VmInfo",
+    "VmNotFoundError",
+    "VmSize",
+    "VmState",
+]

azure_vm/_backend.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import subprocess
+from dataclasses import dataclass
+from typing import Protocol
+
+
+@dataclass
+class CommandResult:
+    args: list[str]
+    returncode: int
+    stdout: str
+    stderr: str
+
+    @property
+    def success(self) -> bool:
+        return self.returncode == 0
+
+
+class CommandBackend(Protocol):
+    def run(
+        self,
+        args: list[str],
+        *,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+    ) -> CommandResult: ...
+
+
+class TofuBackend:
+    """Real backend — invokes the OpenTofu CLI via subprocess."""
+
+    def run(
+        self,
+        args: list[str],
+        *,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+    ) -> CommandResult:
+        try:
+            proc = subprocess.run(
+                args, capture_output=True, text=True, cwd=cwd, env=env
+            )
+        except FileNotFoundError as err:
+            from .exceptions import TofuNotInstalledError
+
+            raise TofuNotInstalledError() from err
+        return CommandResult(
+            args=args,
+            returncode=proc.returncode,
+            stdout=proc.stdout,
+            stderr=proc.stderr,
+        )
+
+
+class FakeBackend:
+    """Test backend — returns pre-configured responses and records all calls."""
+
+    def __init__(
+        self,
+        responses: dict[tuple[str, ...], CommandResult] | None = None,
+    ) -> None:
+        self._responses: dict[tuple[str, ...], CommandResult] = responses or {}
+        self._queues: dict[tuple[str, ...], list[CommandResult]] = {}
+        self._calls: list[list[str]] = []
+        self._cwds: list[str | None] = []
+        self._envs: list[dict[str, str] | None] = []
+        self._default: CommandResult | None = None
+
+    def set_default(self, result: CommandResult) -> None:
+        self._default = result
+
+    def push(self, *args: str, result: CommandResult) -> None:
+        self._queues.setdefault(args, []).append(result)
+
+    def run(
+        self,
+        args: list[str],
+        *,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+    ) -> CommandResult:
+        self._calls.append(list(args))
+        self._cwds.append(cwd)
+        self._envs.append(env)
+        key = tuple(args)
+        if self._queues.get(key):
+            return self._queues[key].pop(0)
+        if key in self._responses:
+            return self._responses[key]
+        if self._default is not None:
+            return self._default
+        raise KeyError(f"FakeBackend: no response configured for {args!r}")
+
+    @property
+    def calls(self) -> list[list[str]]:
+        return list(self._calls)
+
+    @property
+    def cwds(self) -> list[str | None]:
+        return list(self._cwds)
+
+    def last_call(self) -> list[str]:
+        return self._calls[-1] if self._calls else []
+
+    def last_cwd(self) -> str | None:
+        return self._cwds[-1] if self._cwds else None
+
+    @property
+    def envs(self) -> list[dict[str, str] | None]:
+        return list(self._envs)
+
+    def last_env(self) -> dict[str, str] | None:
+        return self._envs[-1] if self._envs else None
+
+
+def run_command(
+    backend: CommandBackend,
+    args: list[str],
+    *,
+    cwd: str | None = None,
+) -> CommandResult:
+    """Run *args* through *backend* and raise on failure."""
+    from .exceptions import AzureVmCommandError
+
+    result = backend.run(args, cwd=cwd)
+    if not result.success:
+        raise AzureVmCommandError(
+            result.args, result.returncode, result.stdout, result.stderr
+        )
+    return result

azure_vm/_discovery.py

@@ -0,0 +1,50 @@
+"""Azure Marketplace image and VM size discovery."""
+
+from __future__ import annotations
+
+import json
+
+from ._backend import CommandBackend, run_command
+from .models import ImageInfo, VmSize
+
+
+def list_images(
+    backend: CommandBackend,
+    publisher: str = "Canonical",
+) -> list[ImageInfo]:
+    """List available VM images from Azure Marketplace."""
+    result = run_command(
+        backend,
+        [
+            "az",
+            "vm",
+            "image",
+            "list",
+            "--publisher",
+            publisher,
+            "--all",
+            "--output",
+            "json",
+        ],
+    )
+    return ImageInfo.from_az_image_list(json.loads(result.stdout))
+
+
+def list_sizes(
+    backend: CommandBackend,
+    location: str,
+) -> list[VmSize]:
+    """List available VM sizes in an Azure region."""
+    result = run_command(
+        backend,
+        [
+            "az",
+            "vm",
+            "list-sizes",
+            "--location",
+            location,
+            "--output",
+            "json",
+        ],
+    )
+    return VmSize.from_az_vm_size_list(json.loads(result.stdout))

azure_vm/_templates.py

@@ -0,0 +1,238 @@
+"""HCL templates and helper functions for OpenTofu workspace generation."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+
+SHARED_TEMPLATE = """\
+terraform {{
+  required_providers {{
+    azurerm = {{
+      source  = "hashicorp/azurerm"
+      version = "~> 4.0"
+    }}
+  }}
+}}
+
+provider "azurerm" {{
+  features {{}}
+  resource_provider_registrations = "none"
+}}
+
+resource "azurerm_resource_group" "main" {{
+  name     = "{resource_group}"
+  location = "{location}"
+}}
+"""
+
+VM_TEMPLATE = """\
+terraform {{
+  required_providers {{
+    azurerm = {{
+      source  = "hashicorp/azurerm"
+      version = "~> 4.0"
+    }}
+  }}
+}}
+
+provider "azurerm" {{
+  features {{}}
+  resource_provider_registrations = "none"
+}}
+
+data "azurerm_resource_group" "main" {{
+  name = "{resource_group}"
+}}
+
+resource "azurerm_virtual_network" "vm" {{
+  name                = "${{var.vm_name}}-vnet"
+  address_space       = ["10.0.0.0/16"]
+  location            = data.azurerm_resource_group.main.location
+  resource_group_name = data.azurerm_resource_group.main.name
+}}
+
+resource "azurerm_subnet" "vm" {{
+  name                 = "${{var.vm_name}}-subnet"
+  resource_group_name  = data.azurerm_resource_group.main.name
+  virtual_network_name = azurerm_virtual_network.vm.name
+  address_prefixes     = ["10.0.1.0/24"]
+}}
+
+resource "azurerm_network_security_group" "vm" {{
+  name                = "${{var.vm_name}}-nsg"
+  location            = data.azurerm_resource_group.main.location
+  resource_group_name = data.azurerm_resource_group.main.name
+
+  security_rule {{
+    name                       = "SSH"
+    priority                   = 1001
+    direction                  = "Inbound"
+    access                     = "Allow"
+    protocol                   = "Tcp"
+    source_port_range          = "*"
+    destination_port_range     = "22"
+    source_address_prefix      = "*"
+    destination_address_prefix = "*"
+  }}
+{extra_security_rules}}}
+
+resource "azurerm_subnet_network_security_group_association" "vm" {{
+  subnet_id                 = azurerm_subnet.vm.id
+  network_security_group_id = azurerm_network_security_group.vm.id
+}}
+
+resource "azurerm_public_ip" "vm" {{
+  name                = "${{var.vm_name}}-pip"
+  location            = data.azurerm_resource_group.main.location
+  resource_group_name = data.azurerm_resource_group.main.name
+  allocation_method   = "Static"
+}}
+
+resource "azurerm_network_interface" "vm" {{
+  name                = "${{var.vm_name}}-nic"
+  location            = data.azurerm_resource_group.main.location
+  resource_group_name = data.azurerm_resource_group.main.name
+
+  ip_configuration {{
+    name                          = "internal"
+    subnet_id                     = azurerm_subnet.vm.id
+    private_ip_address_allocation = "Dynamic"
+    public_ip_address_id          = azurerm_public_ip.vm.id
+  }}
+}}
+
+resource "azurerm_linux_virtual_machine" "vm" {{
+  name                  = var.vm_name
+  location              = data.azurerm_resource_group.main.location
+  resource_group_name   = data.azurerm_resource_group.main.name
+  network_interface_ids = [azurerm_network_interface.vm.id]
+  size                  = var.vm_size
+  computer_name         = var.vm_name
+  admin_username        = "{ssh_username}"
+  disable_password_authentication = true
+
+  admin_ssh_key {{
+    username   = "{ssh_username}"
+    public_key = file("{ssh_public_key_path}")
+  }}
+
+  os_disk {{
+    caching              = "ReadWrite"
+    storage_account_type = "Standard_LRS"
+    disk_size_gb         = var.disk_size_gb
+  }}
+
+  source_image_reference {{
+    publisher = "{image_publisher}"
+    offer     = "{image_offer}"
+    sku       = "{image_sku}"
+    version   = "{image_version}"
+  }}
+{custom_data_block}
+}}
+
+variable "vm_name" {{
+  type = string
+}}
+
+variable "vm_size" {{
+  type    = string
+  default = "Standard_B1s"
+}}
+
+variable "disk_size_gb" {{
+  type    = number
+  default = 30
+}}
+
+variable "desired_state" {{
+  type    = string
+  default = "running"
+}}
+
+output "vm_ip" {{
+  value = azurerm_public_ip.vm.ip_address
+}}
+
+output "vm_state" {{
+  value = var.desired_state
+}}
+
+output "location" {{
+  value = data.azurerm_resource_group.main.location
+}}
+
+output "vm_size" {{
+  value = var.vm_size
+}}
+
+output "image_urn" {{
+  value = "{image_urn}"
+}}
+
+output "resource_group" {{
+  value = "{resource_group}"
+}}
+"""
+
+
+def parse_image_urn(image_urn: str | None) -> tuple[str, str, str, str]:
+    """Parse an Azure image URN into (publisher, offer, sku, version)."""
+    if not image_urn:
+        return ("Canonical", "ubuntu-24_04-lts", "server-gen1", "latest")
+    parts = image_urn.split(":")
+    return (
+        parts[0] if len(parts) > 0 and parts[0] else "Canonical",
+        parts[1] if len(parts) > 1 and parts[1] else "ubuntu-24_04-lts",
+        parts[2] if len(parts) > 2 and parts[2] else "server-gen1",
+        parts[3] if len(parts) > 3 and parts[3] else "latest",
+    )
+
+
+def resolve_ssh_path(
+    ssh_key_path: str | None,
+    default_path: str | None = None,
+) -> str:
+    """Resolve an SSH key path, expanding ~ and applying fallback chain."""
+    return str(Path(ssh_key_path or default_path or "~/.ssh/id_rsa.pub").expanduser())
+
+
+def write_cloud_init(workspace: Path, config: dict | str | None) -> str:
+    """Write cloud-init file to *workspace*. Returns the HCL custom_data block."""
+    if config is None:
+        return ""
+    if isinstance(config, dict):
+        content = "#cloud-config\n" + yaml.dump(config, default_flow_style=False)
+    else:
+        content = config
+    (workspace / "cloud-init.yaml").write_text(content)
+    return '  custom_data = filebase64("cloud-init.yaml")\n'
+
+
+def render_security_rules(open_ports) -> str:
+    """Extra inbound TCP allow-rules for the per-VM NSG.
+
+    Port 22 is skipped (the SSH rule is always rendered). Source is "*",
+    matching the SSH rule's posture — these are short-lived e2e VMs; source
+    restriction is a possible future enhancement.
+    """
+    rules: list[str] = []
+    for index, port in enumerate(sorted({int(p) for p in (open_ports or [])})):
+        if port == 22:
+            continue
+        rules.append(
+            "\n  security_rule {\n"
+            f'    name                       = "Port{port}"\n'
+            f"    priority                   = {1010 + index}\n"
+            '    direction                  = "Inbound"\n'
+            '    access                     = "Allow"\n'
+            '    protocol                   = "Tcp"\n'
+            '    source_port_range          = "*"\n'
+            f'    destination_port_range     = "{port}"\n'
+            '    source_address_prefix      = "*"\n'
+            '    destination_address_prefix = "*"\n'
+            "  }\n"
+        )
+    return "".join(rules)

azure_vm/_workspace.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from pathlib import Path
+
+from ._backend import CommandBackend, run_command
+from ._templates import (
+    SHARED_TEMPLATE,
+    VM_TEMPLATE,
+    parse_image_urn,
+    render_security_rules,
+    resolve_ssh_path,
+    write_cloud_init,
+)
+from .models import VmInfo
+
+
+class Workspace:
+    """Encapsulates filesystem operations on the Azure VM workspace directory.
+
+    Each VM gets a sub-directory under *root*. Shared infrastructure lives in
+    ``.shared/``.
+    """
+
+    def __init__(self, root: Path, backend: CommandBackend) -> None:
+        self._root = root
+        self._backend = backend
+
+    # ------------------------------------------------------------------- paths
+
+    def vm_dir(self, name: str) -> Path:
+        return self._root / name
+
+    def shared_dir(self) -> Path:
+        return self._root / ".shared"
+
+    def vm_exists(self, name: str) -> bool:
+        return (self.vm_dir(name) / "main.tf").exists()
+
+    # ----------------------------------------------------------- shared infra
+
+    def ensure_shared_infra(self, resource_group: str, location: str) -> None:
+        shared = self.shared_dir()
+        hcl = SHARED_TEMPLATE.format(
+            resource_group=resource_group,
+            location=location,
+        )
+        main_tf = shared / "main.tf"
+        # Re-render from the current configuration: an existing workspace must
+        # never silently pin stale parameters. Re-apply only on actual change.
+        if main_tf.exists() and main_tf.read_text() == hcl:
+            return
+        shared.mkdir(parents=True, exist_ok=True)
+        main_tf.write_text(hcl)
+        run_command(self._backend, ["tofu", "init"], cwd=str(shared))
+        _try_import_resource_group(self._backend, shared, resource_group)
+        run_command(self._backend, ["tofu", "apply", "-auto-approve"], cwd=str(shared))
+
+    # --------------------------------------------------------------- vm ops
+
+    def write_vm_workspace(
+        self,
+        name: str,
+        vm_size: str,
+        disk_size_gb: int,
+        image_urn: str | None,
+        cloud_init_config: dict | str | None,
+        ssh_key_path: str | None,
+        *,
+        resource_group: str,
+        location: str,
+        ssh_username: str,
+        default_ssh_key: str | None = None,
+        open_ports=None,
+    ) -> None:
+        workspace = self.vm_dir(name)
+        workspace.mkdir(parents=True, exist_ok=True)
+
+        cloud_init_block = write_cloud_init(workspace, cloud_init_config)
+        publisher, offer, sku, version = parse_image_urn(image_urn)
+        effective_ssh_path = resolve_ssh_path(ssh_key_path, default_ssh_key)
+
+        full_urn = image_urn or "Canonical:ubuntu-24_04-lts:server-gen1:latest"
+
+        hcl = VM_TEMPLATE.format(
+            resource_group=resource_group,
+            location=location,
+            ssh_username=ssh_username,
+            ssh_public_key_path=effective_ssh_path,
+            image_publisher=publisher,
+            image_offer=offer,
+            image_sku=sku,
+            image_version=version,
+            image_urn=full_urn,
+            custom_data_block=cloud_init_block,
+            extra_security_rules=render_security_rules(open_ports),
+        )
+        (workspace / "main.tf").write_text(hcl)
+
+        tfvars = (
+            f'vm_name = "{name}"\n'
+            f'vm_size = "{vm_size}"\n'
+            f"disk_size_gb = {disk_size_gb}\n"
+        )
+        (workspace / "terraform.tfvars").write_text(tfvars)
+
+    def iter_vm_workspaces(self) -> Iterator[Path]:
+        """Yield VM workspace directories in sorted order.
+
+        Hidden entries, non-directories, and directories that hold no
+        ``main.tf`` are skipped.
+        """
+        for item in sorted(self._root.iterdir()):
+            if item.name.startswith("."):
+                continue
+            if not item.is_dir():
+                continue
+            if not (item / "main.tf").exists():
+                continue
+            yield item
+
+    def list_vms(self) -> list[VmInfo]:
+        vms: list[VmInfo] = []
+        for vm_dir_item in self.iter_vm_workspaces():
+            result = self._backend.run(
+                ["tofu", "output", "-json"], cwd=str(vm_dir_item)
+            )
+            if result.success and result.stdout:
+                try:
+                    data = json.loads(result.stdout)
+                    vms.append(VmInfo.from_tofu_output(data, vm_dir_item.name))
+                except json.JSONDecodeError:
+                    continue
+        return vms
+
+    def purge(self) -> None:
+        """Destroy all VM workspaces (preserving .shared/)."""
+        for vm_dir_item in self.iter_vm_workspaces():
+            run_command(
+                self._backend,
+                ["tofu", "destroy", "-auto-approve"],
+                cwd=str(vm_dir_item),
+            )
+
+
+def _try_import_resource_group(
+    backend: CommandBackend,
+    shared: Path,
+    resource_group: str,
+) -> None:
+    """Try to import an existing Azure resource group into Tofu state.
+
+    If the resource group already exists in Azure but is not tracked in
+    the local Tofu state (e.g. fresh workspace, or state was deleted),
+    ``tofu apply`` would fail with "already exists". Importing it first
+    makes apply a no-op.  When the import fails the resource group
+    likely does not exist yet — apply will create it.
+    """
+    # Resolve the subscription id via Azure CLI (must be authenticated).
+    result = backend.run(
+        ["az", "account", "show", "--query", "id", "-o", "tsv"],
+    )
+    if not result.success:
+        return  # can't query; let apply try anyway
+    sub_id = (result.stdout or "").strip()
+    if not sub_id:
+        return
+
+    resource_id = f"/subscriptions/{sub_id}/resourceGroups/{resource_group}"
+    # Intentionally ignore result — import is best-effort.
+    backend.run(
+        ["tofu", "import", "azurerm_resource_group.main", resource_id],
+        cwd=str(shared),
+    )