Source code for pyintrovirt.domain

"""Classes for dealing with a domain/VM being introspected."""

from contextlib import ContextDecorator
from typing import NamedTuple, Optional, Union

import introvirt  # pylint: disable=import-error




class DomainInformation(NamedTuple):
    """Domain information."""

    domain_name: str
    domain_id: int

    def __str__(self):
        return f"DomainInformation(domain_name={self.domain_name}, domain_id={self.domain_id})"





class Domain(ContextDecorator):
    """An introspection target (VM/domain) that we are attached to for introspecting."""

    def __init__(self, domain_id: Union[int, str], hypervisor: introvirt.Hypervisor):
        """
        Initialize an target domain to be introspected.

        Args:
            domain_id: The domain to attach to. Can be an integer domain ID or a string domain name.
            hypervisor: The hypervisor instance with the target domain to attach to.
        """
        domain = hypervisor.attach_domain(domain_id)
        if not domain.detect_guest():
            raise RuntimeError("Failed to detect guest OS")
        self._domain: Optional[introvirt.Domain] = domain
        self._guest: introvirt.Guest = domain.guest()
        self._os: introvirt.OS = introvirt.OS(self._guest.os())

    def _attached(self) -> introvirt.Domain:
        if self._domain is None:
            raise RuntimeError("Domain is detached")
        return self._domain

    def __enter__(self):
        return self

    def __exit__(self, _exc_type, _exc_value, _traceback):
        self.detach()

    def __del__(self):
        self.detach()

    @property
    def os(self) -> introvirt.OS:
        """Access the guest OS type enum."""
        return self._os

    def _windows_guest(self) -> introvirt.WindowsGuest:
        win_guest = introvirt.WindowsGuest_from_guest(self._guest)
        if win_guest is None:
            raise RuntimeError("Guest is not a Windows guest")
        return win_guest

    @property
    def syscall_categories(self) -> tuple[str, ...]:
        """Get a list of system call categories."""
        if self.os == introvirt.OS.Windows:
            return self._windows_guest().syscall_categories()  # ty: ignore[too-many-positional-arguments]
        raise NotImplementedError("Only implemented for Windows guests right now.")



    def detach(self) -> None:
        """Detach from the domain. Safe to call multiple times."""
        if self._domain is None:
            return
        self._domain.interrupt()
        self._domain = None




    def poll(self, event_handler: introvirt.EventCallback):
        """Start the poller for events. No events will be recieved until this is started."""
        self._attached().poll(event_handler)




    def clear_system_call_filter(self):
        """Clear the system call filter if set."""
        self._attached().system_call_filter().clear()




    def clear_task_filter(self):
        """Clear the process filter if set."""
        self._attached().task_filter().clear()




    def default_system_call_filter(self):
        """Set the system call filter to the default set of supported system calls for the OS."""
        if self.os == introvirt.OS.Windows:
            self._windows_guest().default_syscall_filter(self._attached().system_call_filter())
        else:
            raise NotImplementedError("Only implemented for Windows guests right now.")




    def filter_system_call(self, syscall: introvirt.SystemCallIndex, enabled: bool):
        """Toggle filtering of a specific system call."""
        if self.os == introvirt.OS.Windows:
            self._windows_guest().set_system_call_filter(self._attached().system_call_filter(), syscall.value, enabled)
        else:
            raise NotImplementedError("Only implemented for Windows guests right now.")




    def filter_system_call_category(self, category: str):
        """Filter by a system call category."""
        if self.os == introvirt.OS.Windows:
            self._windows_guest().enable_category(category, self._attached().system_call_filter())
        else:
            raise NotImplementedError("Only implemented for Windows guests right now.")




    def filter_system_calls(self, enabled: bool):
        """Toggle system call filtering on/off. Required for the filter to take effect."""
        self._attached().system_call_filter().enabled(enabled)




    def filter_task_name(self, name: str):
        """Filter by task name. Name is a case-insensitive process name prefix."""
        self._attached().task_filter().add_name(name)




    def filter_task_pid(self, pid: int):
        """Filter by task PID."""
        self._attached().task_filter().add_pid(pid)




    def filter_task_tid(self, tid: int):
        """Filter by task TID."""
        self._attached().task_filter().add_tid(tid)




    def unfilter_task_name(self, name: str):
        """Remove a name from the task filter."""
        self._attached().task_filter().remove_name(name)




    def unfilter_task_pid(self, pid: int):
        """Remove a PID from the task filter."""
        self._attached().task_filter().remove_pid(pid)




    def unfilter_task_tid(self, tid: int):
        """Remove a TID from the task filter."""
        self._attached().task_filter().remove_tid(tid)




    def intercept_system_calls(self, enabled: bool):
        """Toggle system call interception on/off. Required to received system call events at all."""
        self._attached().intercept_system_calls(enabled)




    def intercept_cr_writes(self, cr: int, enabled: bool):
        """Intercept CR writes."""
        self._attached().intercept_cr_writes(cr, enabled)