Mini Shell

"""VZ / HA functions"""

import enum
import json
import os
import shlex
import subprocess

from pathlib import Path
from typing import Iterable, Optional, Union

import distro


class CT:
    """Create a CT object representing a VZ Container"""

    def __init__(self, ctid: str):
        try:
            self.ctid: str = str(ctid)
        except ValueError as exc:
            raise VZError(f"Provided ctid '{ctid}' is not of type str()") from exc

        # List of options for vzlist
        opts = [
            "cpulimit",
            "cpus",
            "cpuunits",
            "ctid",
            "description",
            "device",
            "disabled",
            "diskinodes",
            "diskspace",
            "hostname",
            "ip",
            "laverage",
            "name",
            "netfilter",
            "numiptent",
            "numproc",
            "onboot",
            "ostemplate",
            "physpages",
            "private",
            "root",
            "status",
            "swappages",
            "uptime",
            "uuid",
            "veid",
        ]

        # Build vzlist command
        command = ["vzlist", "-jo", ",".join(opts), self.ctid]

        # Execute vzlist command
        result = subprocess.run(command, capture_output=True, text=True, check=False)

        # Parse vzlist command result
        if result.returncode:
            raise VZError(f"{result.stderr}")

        data = json.loads(result.stdout)[0]

        self.cpulimit: int = data.get("cpulimit", 0)
        self.cpus: int = data.get("cpus", 0)
        self.cpuunits: int = data.get("cpuunits", 0)
        self.ctid: Union[int, str] = data.get("ctid", "")
        self.description: str = data.get("description", "")
        self.device: str = data.get("device", "")
        self.disabled: bool = data.get("disabled", False)
        self.diskinodes: dict = data.get("diskinodes", {})
        self.diskspace: dict = data.get("diskspace", {})
        self.hostname: str = data.get("hostname", "")
        self.ip: list[str] = data.get("ip", [])
        self.laverage: list[float] = data.get("laverage", [])
        self.name: str = data.get("name", "")
        self.netfilter: str = data.get("netfilter", "")
        self.numiptent: dict = data.get("numiptent", {})
        self.numproc: dict = data.get("numproc", {})
        self.onboot: bool = data.get("onboot", False)
        self.ostemplate: str = data.get("ostemplate", "")
        self.physpages: dict = data.get("physpages", {})
        self.private: str = data.get("private", "")
        self.root: str = data.get("root", "")
        self.status: str = data.get("status", "")
        self.swappages: dict = data.get("swappages", {})
        self.uptime: float = data.get("uptime", 0.0)
        self.uuid: str = data.get("uuid", "")
        self.veid: str = data.get("veid", "")

    def __repr__(self):
        attrs = ", ".join(f"{key}={repr(getattr(self, key))}" for key in self.__dict__)
        return f"CT({attrs})"

    def start(self) -> subprocess.CompletedProcess:
        """
        Starts the container identified by self.ctid using the 'vzctl start' command.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl start' command execution.
        """
        command = ["vzctl", "start", self.ctid]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def stop(self) -> subprocess.CompletedProcess:
        """
        Stops the container identified by self.ctid using the 'vzctl stop' command.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl stop' command execution.
        """
        command = ["vzctl", "stop", self.ctid]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def restart(self) -> subprocess.CompletedProcess:
        """
        Restarts the container identified by self.ctid using the 'vzctl restart' command.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl restart' command execution.
        """
        command = ["vzctl", "restart", self.ctid]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def mount(self) -> subprocess.CompletedProcess:
        """
        Mounts the container's filesystem identified by self.ctid using the 'vzctl mount' command.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl mount' command execution.
        """
        command = ["vzctl", "mount", self.ctid]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def umount(self) -> subprocess.CompletedProcess:
        """
        Unmounts the container's filesystem identified by self.ctid using the 'vzctl umount' command.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl umount' command execution.
        """
        command = ["vzctl", "umount", self.ctid]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def destroy(self) -> subprocess.CompletedProcess:
        """
        Destroys the container identified by self.ctid using the 'vzctl destroy' command.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl destroy' command execution.
        """
        command = ["vzctl", "destroy", self.ctid]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def register(self, dst) -> subprocess.CompletedProcess:
        """
        Registers the container at the specified destination using the 'vzctl register' command.

        Args:
            dst (str): The destination path where the container is to be registered.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl register' command execution.
        """
        command = ["vzctl", "register", dst, self.ctid]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def unregister(self) -> subprocess.CompletedProcess:
        """
        Unregisters the container identified by self.ctid using the 'vzctl unregister' command.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl unregister' command execution.
        """
        command = ["vzctl", "unregister", self.ctid]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def clone(self, name: str) -> subprocess.CompletedProcess:
        """
        Clones the container identified by self.ctid with the specified name using the 'prlctl clone' command.

        Args:
            name (str): The name for the new cloned container.

        Returns:
            subprocess.CompletedProcess: The result of the 'prlctl clone' command execution.
        """
        command = ["prlctl", "clone", self.ctid, "--name", str(name)]
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def exec2(self, cmd: str) -> subprocess.CompletedProcess:
        """
        Executes the specified command within the container identified by self.ctid using the 'vzctl exec2' command.

        Args:
            cmd (str): The command to execute within the container.

        Returns:
            subprocess.CompletedProcess: The result of the 'vzctl exec2' command execution.
        """
        command = ["vzctl", "exec2", self.ctid] + str(cmd).split()
        return subprocess.run(command, capture_output=True, text=True, check=False)

    def set(self, **kwargs) -> subprocess.CompletedProcess:
        """
        Sets various configuration options for the container identified by self.ctid.

        This method constructs a command to modify the container's configuration using
        the `vzctl set` command. It accepts keyword arguments where each key-value pair
        represents an option and its corresponding value to be set. The method appends
        "--save" to the command to save the changes permanently. After running the
        command, the method reinitializes the object to reflect the updated configuration.

        Args:
            **kwargs: Arbitrary keyword arguments representing configuration options and values.

        Returns:
            subprocess.CompletedProcess: The result of the `subprocess.run` call,
            containing information about the execution of the command.
        """
        arg_list = []

        for key, value in kwargs.items():
            arg_list.append(f"--{key}")
            arg_list.append(f"{value}")
        arg_list.append("--save")

        command = ["vzctl", "set", self.ctid] + arg_list
        result = subprocess.run(command, capture_output=True, text=True, check=False)

        self.__init__(self.ctid)

        return result

    def fix_ctid(self, start=True) -> None:
        """
        A tool for offline fixing CTID fields of VZ7 containers. Re-registers the container while offline.

        This method stops the container, unregisters it, renames the container directory from the current CTID
        to the container name, updates the instance attributes accordingly, and re-registers the container with
        the new details. Optionally, it can restart the container if the `start` parameter is set to True.

        Args:
            start (bool): If True, the container will be restarted after the CTID fix is applied. Defaults to True.

        Returns:
            None
        """
        self.stop()
        self.unregister()

        src = Path(f"/vz/private/{self.ctid}")
        dst = Path(f"/vz/private/{self.name}")
        src.rename(dst)

        self.private = str(dst)
        self.root = f"/vz/root/{self.name}"
        self.ctid = self.veid = self.name

        self.register(self.private)

        if start:
            self.start()


class ListCmd(enum.Enum):
    """vz list base commands"""

    VZLIST = ["/usr/sbin/vzlist", "-H"]
    PRLCTL = ["/usr/bin/prlctl", "list", "-H"]


class VZError(Exception):
    """Raised for errors with VZ and OpenVZ"""


def is_vz() -> bool:
    """Checks if host is a Virtuozzo node"""
    return bool("Virtuozzo" in distro.name())


def is_openvz() -> bool:
    """Check if host is an OpenVZ node"""
    return os.path.isfile("/etc/virtuozzo-release") and not is_vz()


def is_vz7() -> bool:
    """Check if host is a Virtuozzo 7 node"""
    return bool(is_vz() and distro.major_version() == "7")


def is_vps() -> bool:
    """Check if host is a Virtuozzo container"""
    try:
        with open("/proc/vz/veinfo", encoding="ascii") as handle:
            ve_data = handle.read().strip()
    except IOError:
        return False  # if veinfo doesn't exist this can't be a vps
    if ve_data.count("\n") != 0:
        return False
    try:
        veid = int(
            ve_data.split()[0]
        )  # if veinfo contains >1 line, this is a CL or VZ node
    except ValueError:
        return True  # veinfo contains a UUID
    return veid != 0


def _exec(cmd: Iterable):
    """For executing prlctl or vzlist"""
    try:
        ret = subprocess.run(cmd, capture_output=True, encoding="utf-8", check=False)
    except FileNotFoundError as exc:
        raise VZError(exc) from exc
    if ret.returncode:  # nonzero
        raise VZError(f"Error running {cmd!r}. stderr={ret.stderr!r}")
    return ret


def is_ct_running(ctid: Union[str, int]) -> bool:
    """Checks if a container is running

    Args:
        ctid: container ID to check
    Returns:
        True if the container is running on this node, False if it
        isn't or if some other error occurs
    """
    try:
        ret = subprocess.run(
            ["/usr/bin/prlctl", "list", "-H", "-o", "status", str(ctid)],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            encoding="utf-8",
            check=True,
        )
    except FileNotFoundError:
        pass  # retry with vzlist
    except subprocess.CalledProcessError:
        return False
    else:
        return ret.stdout.split()[0] == "running"
    try:
        ret = _exec(["/usr/sbin/vzlist", "-H", "-o", "status", str(ctid)])
    except VZError:  # CTID probably doesn't exist
        return False
    return ret.stdout.split()[0] == "running"


def uuid2ctid(uuid: str) -> str:
    """get the legacy CTID of a container

    Args:
        uuid: VZ UUID to find the legacy CTID for
    Raises:
        VZError: if the prlctl command fails
    """
    ret = _exec(["/usr/bin/prlctl", "list", "-H", "-o", "name", uuid])
    return ret.stdout.split()[0]


def ctid2uuid(ctid: Union[int, str]) -> str:
    """Obtain the UUID of a container from its legacy CTID

    Warning:
        This does not work on VZ4
    Args:
        ctid: Legacy CTID to get the UUID for
    Raises:
        VZError: if the prlctl command fails
    """
    ret = _exec(["/usr/bin/prlctl", "list", "-H", "-o", "uuid", str(ctid)])
    return ret.stdout.split()[0].strip(r"{}")


def get_envid(ctid: Union[int, str]) -> str:
    """Obtain the EnvID of a container

    Note:
        This determines what the subdirectory of /vz/root and /vz/private will
        be. This also has to run on VZ4 which lacks the envid field or prlctl,
        so we just return the CTID
    Args:
        ctid: legacy CTID to find the envid for
    Raises:
        VZError: if the prlctl command fails or
            /etc/virtuozzo-release is missing
    """
    try:
        with open("/etc/virtuozzo-release", "r", encoding="utf-8") as handle:
            if "Virtuozzo release 4" in handle.read():
                return str(ctid)
    except FileNotFoundError as exc:
        raise VZError(exc) from exc
    ret = _exec(["/usr/bin/prlctl", "list", "-H", "-o", "envid", str(ctid)])
    return ret.stdout.split()[0]


def _list_cmd(
    opts: list,
    args: Optional[list] = None,
    list_cmd: Optional[ListCmd] = None,
) -> tuple[ListCmd, list[str]]:
    """Deterines the cmd to run based on VZ version for get_cts()

    Args:
        opts: items to send into ``-o/--output``
        args: optional params to send such as ``--all``
        list_cmd (ListCmd): set this to ListCmd.VZLIST or ListCmd.PRLCTL to
            skip auto-detecting which command to use
    """
    if list_cmd is None:
        if is_vz():
            if not is_vz7() and "ostemplate" in opts:
                # prlctl's ostemplate is broken and reports distro on vz6
                # switch to vzlist; fix envid to veid if it was requested
                list_cmd = ListCmd.VZLIST
            else:
                list_cmd = ListCmd.PRLCTL
        else:  # OpenVZ
            list_cmd = ListCmd.VZLIST
    if list_cmd == ListCmd.VZLIST:
        conv_opts = {x: ("veid" if x == "envid" else x) for x in opts}
    else:
        # prctl refers to 'ctid' as 'name'
        conv_opts = {x: ("name" if x == "ctid" else x) for x in opts}
    cmd = list_cmd.value.copy()
    if args is not None:
        cmd.extend(args)
    # forces opts's vals to be in the same order as args
    cmd_opts = ",".join([conv_opts[x] for x in opts])
    cmd.extend(["-o", cmd_opts])
    return list_cmd, cmd


def _read_row(
    list_cmd: ListCmd, cmd: list[str], row: list[str], opts: list[str]
) -> dict[str, str]:
    # if number of rows matches requested options, return normally
    if len(row) == len(opts):
        return {x: row[i] for i, x in enumerate(opts)}
    # handle an edge case: prlctl can print missing ostemplates as '' while
    # vzlist prints it as '-', making the prlctl one harder to parse
    if (
        list_cmd == ListCmd.PRLCTL
        and len(row) == len(opts) - 1
        and "ostemplate" in opts
    ):
        opts = opts.copy()
        opts.remove("ostemplate")
        ret = {x: row[i] for i, x in enumerate(opts)}
        ret["ostemplate"] = "-"
        return ret
    raise VZError(
        f"{shlex.join(cmd)} expected {len(opts)} columns," f" but got {len(row)}: {row}"
    )


def get_cts(
    opts: Optional[list] = None,
    args: Optional[list] = None,
    list_cmd: Optional[ListCmd] = None,
) -> list[dict[str, str]]:
    """Returns containers according to platform as a list of dicts

    Args:
        opts: items to send into -o/--output (will default to ['ctid'] if None)
        args: optional params to send such as --all
        list_cmd (ListCmd): set this to ListCmd.VZLIST or ListCmd.PRLCTL to
            skip auto-detecting which command to use
    Raises:
        VZError: if the prlctl or vzlist command fails
    """
    if not opts:
        opts = ["ctid"]
    ret = []
    # process each line as a dict where keys are the arg and vals are the result
    list_cmd, cmd = _list_cmd(opts, args, list_cmd)
    for row in _exec(cmd).stdout.splitlines():
        row = row.strip()
        if not row:
            continue  # blank line
        ret.append(_read_row(list_cmd, cmd, row.split(), opts))
    return ret