"""Device-specific fix executors.

- :func:`device_swap_vhost_to_qemu` — flip every NIC's
  ``<driver name='vhost'/>`` to ``name='qemu'`` (userspace
  virtio-net; smaller kernel-CVE blast radius).
- :func:`device_strip_tpm_passthrough` — remove only ``<tpm>``
  elements whose backend is ``passthrough``; emulated TPMs are
  fine.
"""

from __future__ import annotations

import xml.etree.ElementTree as ET  # noqa: TC003
from typing import TYPE_CHECKING, Any

from mgf.vm.core.security.isolation.fixes._helpers import devices_element
from mgf.vm.core.security.isolation.types import isolation_fix

if TYPE_CHECKING:
    from mgf.vm.core.common.progress import ProgressReporter


__all__ = [
    "device_strip_tpm_passthrough",
    "device_swap_vhost_to_qemu",
]


@isolation_fix("device.swap_vhost_to_qemu")
def device_swap_vhost_to_qemu(
    *,
    target: ET.ElementTree,
    reporter: ProgressReporter,
    dry_run: bool = False,
    **_kw: Any,
) -> None:
    """Set every NIC ``<driver name>`` from ``vhost`` to ``qemu``."""
    devices = devices_element(target)
    if devices is None:
        reporter.report("device.swap_vhost_to_qemu: no <devices> element")
        return

    changed = 0
    for iface in devices.findall("interface"):
        driver = iface.find("driver")
        if driver is None:
            continue
        if driver.get("name") != "vhost":
            continue
        if dry_run:
            reporter.report(
                "device.swap_vhost_to_qemu: (dry-run) would set name='qemu'",
            )
        else:
            driver.set("name", "qemu")
        changed += 1

    if changed == 0:
        reporter.report("device.swap_vhost_to_qemu: already satisfied")
        return
    if not dry_run:
        reporter.report(
            f"device.swap_vhost_to_qemu: rewrote {changed} interface driver(s)",
        )


@isolation_fix("device.strip_tpm_passthrough")
def device_strip_tpm_passthrough(
    *,
    target: ET.ElementTree,
    reporter: ProgressReporter,
    dry_run: bool = False,
    **_kw: Any,
) -> None:
    """Remove only ``<tpm>`` elements whose backend is ``passthrough``.

    Emulated TPMs (``<backend type='emulator'/>``) are fine and
    are left in place.
    """
    devices = devices_element(target)
    if devices is None:
        reporter.report("device.strip_tpm_passthrough: no <devices> element")
        return

    bad: list[ET.Element] = []
    for tpm in devices.findall("tpm"):
        backend = tpm.find("backend")
        if backend is not None and backend.get("type") == "passthrough":
            bad.append(tpm)

    if not bad:
        reporter.report("device.strip_tpm_passthrough: already satisfied")
        return
    if dry_run:
        reporter.report(
            f"device.strip_tpm_passthrough: (dry-run) would remove {len(bad)} passthrough TPM(s)",
        )
        return
    for tpm in bad:
        devices.remove(tpm)
    reporter.report(
        f"device.strip_tpm_passthrough: removed {len(bad)} passthrough TPM(s)",
    )
