common.py

from __future__ import annotations

import pytest

import getpass
import inspect
import itertools
import logging
import multiprocessing
import os
import random
import string
import sys
import tempfile
import time
import traceback
from datetime import datetime
from enum import Enum
from functools import lru_cache
from pathlib import Path
from uuid import UUID

import requests
from passlib.hash import sha512_crypt
from pydantic import TypeAdapter

from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Literal,
    TypeAlias,
    TypeVar,
    cast,
    overload,
)

if TYPE_CHECKING:
    from lib.host import Host


KiB = 2**10
MiB = KiB**2
GiB = KiB**3
TiB = KiB**4

VHD_MAX = 2040 * GiB
QCOW2_MAX = 16 * TiB - 2561 * MiB

_SYMBOLIC_SIZES: dict[str, int] = {
    'VHD_MAX': VHD_MAX,
    'QCOW2_MAX': QCOW2_MAX,
}

def parse_size(size_str: str) -> int:
    """
    Parse a size string like "2.5TiB", "1GiB", "1024", "VHD_MAX", or "QCOW2_MAX".
    """
    symbolic = _SYMBOLIC_SIZES.get(size_str.strip().upper())
    if symbolic is not None:
        return symbolic
    try:
        return int(size_str)
    except ValueError:
        pass

    size_str = size_str.strip()
    for unit, multiplier in [('TiB', TiB), ('GiB', GiB), ('MiB', MiB), ('KiB', KiB)]:
        if size_str.endswith(unit):
            try:
                return int(float(size_str[:-len(unit)].strip()) * multiplier)
            except ValueError:
                pass

    raise ValueError(f"Cannot parse size: {size_str}")

T = TypeVar("T")

HostAddress: TypeAlias = str
DiskDevName: TypeAlias = str
Defer: TypeAlias = Callable[[Callable[[], object]], None]

class PackageManagerEnum(Enum):
    UNKNOWN = 1
    YUM = 2
    APT_GET = 3
    APK = 4
    DNF = 5
    ZYPPER = 6

# Common VM images used in tests
def vm_image(vm_key: str) -> str:
    from data import DEF_VM_URL, VM_IMAGES
    url = VM_IMAGES[vm_key]
    if not url.startswith('http'):
        url = DEF_VM_URL + url
    return url

def prefix_object_name(label: str) -> str:
    name_prefix = None
    try:
        from data import OBJECTS_NAME_PREFIX
        name_prefix = OBJECTS_NAME_PREFIX
    except ImportError:
        pass
    if name_prefix is None:
        name_prefix = f"[{getpass.getuser()}]"
    return f"{name_prefix} {label}"

def shortened_nodeid(nodeid: str) -> str:
    components = nodeid.split("::")
    # module
    components[0] = strip_prefix(components[0], "tests/")
    components[0] = strip_suffix(components[0], ".py")
    components[0] = components[0].replace("/", ".")
    # function
    components[-1] = strip_prefix(components[-1], "test_")
    # class
    if len(components) > 2:
        components[1] = strip_prefix(components[1], "Test")

    return "::".join(components)

def expand_scope_relative_nodeid(scoped_nodeid: str, scope: str, ref_nodeid: str) -> str:
    if scope == 'session' or scope == 'package':
        base = []
    elif scope == 'module':
        base = ref_nodeid.split("::", 1)[:1]
    elif scope == 'class':
        base = ref_nodeid.split("::", 2)[:2]
    else:
        raise RuntimeError(f"Internal error: invalid scope {scope!r}")
    logging.debug("scope: %r base: %r relative: %r", scope, base, scoped_nodeid)
    return "::".join(itertools.chain(base, (scoped_nodeid,)))

@lru_cache(maxsize=None)
def _get_type_adapter(tp: Any) -> TypeAdapter[Any]:
    return TypeAdapter(tp)

@overload
def ensure_type(tp: type[T], v: object) -> T:
    ...

@overload
def ensure_type(tp: Any, v: object) -> Any:
    ...

def ensure_type(tp: Any, v: object) -> Any:
    """
    Cast a value to the specified type, and checks it's actual type at run time

    Unlike typing.cast, it also performs a runtime check.
    Unlike isinstance, it also supports complex types.

    Until PEP 747 is accepted and fully implemented, we can't rely on typing_extensions.TypeForm
    to properly infer the type when using Literal, Union, etc. In that case ensure_type() returns Any
    and the value must be explicitly cast to get the actual type in the static type checkers.

    Examples:
    FB = Literal['foo', 'bar']
    x = cast(FB, ensure_type(FB, 'foo'))
    # x is seen as type Literal['foo', 'bar']

    o = object()
    y = ensure_type(int, o)
    # y is seen as type int (without using cast)

    The code for x would validate 'foo' at runtime but would generate an exception for y, because
    o is not actually of type int.
    """
    return _get_type_adapter(tp).validate_python(v)

def callable_marker(value: T | Callable[..., T], request: pytest.FixtureRequest) -> T:
    """
    Process value optionally generated by fixture-dependent callable.

    Typically useful for fixtures using pytest markers on parametrized tests.

    Such markers as parameter one value, or a callable that will
    return a value.  The callable may take as parameters any subset of
    the fixture names the test itself uses.
    """
    if callable(value):
        try:
            params = {arg_name: request.getfixturevalue(arg_name)
                      for arg_name in inspect.getfullargspec(value).args}
        except pytest.FixtureLookupError as e:
            raise RuntimeError("fixture in mapping not found on test") from e
        # callable ensures the value is of type Callable[..., object], which is not enough in that case
        # we can trust the static checker though, and thus use cast
        fn = cast(Callable[..., T], value)
        return fn(**params)
    else:
        return value

def wait_for(fn: Callable[[], object], msg: str | None = None, timeout_secs: int = 2 * 60, retry_delay_secs: int = 2,
             invert: bool = False) -> None:
    if msg is not None:
        logging.info(msg)
    start_time = time.perf_counter()
    while True:
        ret = fn()
        if not invert and ret:
            return
        if invert and not ret:
            return
        if time.perf_counter() - start_time >= timeout_secs:
            expected = 'True' if not invert else 'False'
            raise TimeoutError(
                "Timeout reached while waiting for fn call to yield %s (%s)." % (expected, timeout_secs)
            )
        time.sleep(retry_delay_secs)

def wait_for_not(
    fn: Callable[[], Any], msg: str | None = None, timeout_secs: int = 2 * 60, retry_delay_secs: int = 2
) -> None:
    return wait_for(fn, msg, timeout_secs, retry_delay_secs, True)

def run_with_timeout(fn: Callable[[], Any], timeout_secs: int = 2 * 60) -> None:
    queue: multiprocessing.Queue[Exception] = multiprocessing.Queue()

    def fn_wrapper() -> None:
        try:
            fn()
        except Exception as e:
            queue.put(e)

    proc = multiprocessing.Process(target=fn_wrapper)
    proc.start()
    proc.join(timeout=timeout_secs)

    if proc.is_alive():
        proc.terminate()
        proc.join()

        raise TimeoutError(f"Timeout reached while waiting for fn call to return ({timeout_secs}s).")

    if not queue.empty():
        raise queue.get(block=False)

def is_uuid(maybe_uuid: str) -> bool:
    try:
        UUID(maybe_uuid, version=4)
        return True
    except ValueError:
        return False

def to_xapi_bool(b: bool) -> str:
    return 'true' if b else 'false'

def parse_xe_dict(xe_dict: str) -> dict[str, str]:
    """
    Parses a xe param containing keys and values, e.g. "major: 7; minor: 20; micro: 0; build: 3".

    Data type remains str for all values.
    """
    res = {}
    for pair in xe_dict.split(';'):
        key, value = pair.split(':')
        res[key.strip()] = value.strip()
    return res

def safe_split(text: str, sep: str = ',') -> list[str]:
    """ A split function that returns an empty list if the input string is empty. """
    return text.split(sep) if len(text) > 0 else []

def strip_prefix(string: str, prefix: str) -> str:
    if sys.version_info >= (3, 9):
        return string.removeprefix(prefix)
    if string.startswith(prefix):
        return string[len(prefix):]
    return string

def strip_suffix(string: str, suffix: str) -> str:
    if sys.version_info >= (3, 9):
        return string.removesuffix(suffix)
    if string.endswith(suffix):
        return string[:-len(suffix)]
    return string

def setup_formatted_and_mounted_disk(host: Host, sr_disk: str, fs_type: str, mountpoint: str) -> None:
    if fs_type == 'ext4':
        option_force = '-F'
    elif fs_type == 'xfs':
        option_force = '-f'
    else:
        raise Exception(f"Unsupported fs_type '{fs_type}' in this function")
    device = '/dev/' + sr_disk
    logging.info(f">> Format sr_disk {sr_disk} and mount it on host {host}")
    host.ssh(f'mkfs.{fs_type} {option_force} {device}')
    host.ssh(f'rm -rf {mountpoint}') # Remove any existing leftover to ensure rmdir will not fail in teardown
    host.ssh(f'mkdir -p {mountpoint}')
    host.ssh('cp -f /etc/fstab /etc/fstab.orig')
    ssh_client = host.ssh('echo $SSH_CLIENT').split()[0]
    now = datetime.now().isoformat()
    host.ssh(f'echo "# added by {ssh_client} on {now}\n{device} {mountpoint} {fs_type} defaults 0 0" >>/etc/fstab')
    try:
        host.ssh(f'mount {mountpoint}')
    except Exception:
        # restore fstab then re-raise
        host.ssh('cp -f /etc/fstab.orig /etc/fstab')
        raise

def teardown_formatted_and_mounted_disk(host: Host, mountpoint: str) -> None:
    logging.info(f"<< Restore fstab and unmount {mountpoint} on host {host}")
    host.ssh('cp -f /etc/fstab.orig /etc/fstab')
    host.ssh(f'umount {mountpoint}')
    host.ssh(f'rmdir {mountpoint}')

def exec_nofail(func: Callable[[], Any]) -> list[Exception]:
    """ Execute a function, log a warning if it fails, and return eiter [] or [e] where e is the exception. """
    caller_name = inspect.stack()[1].function
    try:
        func()
        return []
    except Exception as e:
        logging.warning(
            f"An error occurred in `{caller_name}`\n"
            f"Backtrace:\n{traceback.format_exc()}"
        )
        return [e]

def raise_errors(errors: list[Exception]) -> None:
    if not errors:
        return
    elif len(errors) == 1:
        raise errors[0]
    else:
        raise Exception("Several exceptions were catched: " + "\n".join(repr(e) for e in errors))

def strtobool(val: str | None) -> bool:
    # Note: `distutils` package is deprecated and slated for removal in Python 3.12.
    # There is not alternative for strtobool.
    # See: https://peps.python.org/pep-0632/#migration-advice
    # So this is a custom implementation with differences:
    # - A boolean is returned instead of integer
    # - Empty string and None are supported (False is returned in this case)
    if not val:
        return False
    val = val.lower()
    if val in ('y', 'yes', 't', 'true', 'on', '1'):
        return True
    if val in ('n', 'no', 'f', 'false', 'off', '0'):
        return False
    raise ValueError("invalid truth value '{}'".format(val))

def url_download(url: str, filename: str) -> None:
    """
    Download the content of `url` to the `filename` destination.

    A randomized filename is used during download to prevent file corruption on
    concurrent use. If the download fails then the temporary file is removed.
    """
    destination = Path(filename)
    destination.parent.mkdir(parents=True, exist_ok=True)
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        temp_name: str | None = None
        try:
            with tempfile.NamedTemporaryFile(
                dir=destination.parent, prefix=f"{destination.name}.", suffix=".part", delete=False
            ) as fd:
                temp_name = fd.name
                for chunk in r.iter_content(chunk_size=64 * 1024):
                    fd.write(chunk)
        except BaseException:
            if temp_name is not None:
                os.unlink(temp_name)
            raise
        else:
            os.rename(temp_name, filename)

def randid(length: int = 6) -> str:
    """
    Generates a random string of a specified length.
    The string consists of lowercase letters, uppercase letters, and digits.
    """
    characters = string.ascii_lowercase + string.digits
    return ''.join(random.choices(characters, k=length))

@overload
def _param_get(host: Host, xe_prefix: str, uuid: str, param_name: str, key: str | None = ...,
               accept_unknown_key: Literal[False] = ...) -> str:
    ...

@overload
def _param_get(host: Host, xe_prefix: str, uuid: str, param_name: str, key: str | None = ...,
               accept_unknown_key: Literal[True] = ...) -> str | None:
    ...

@overload
def _param_get(host: Host, xe_prefix: str, uuid: str, param_name: str, key: str | None = ...,
               accept_unknown_key: bool = ...) -> str | None:
    ...

def _param_get(host: Host, xe_prefix: str, uuid: str, param_name: str, key: str | None = None,
               accept_unknown_key: bool = False) -> str | None:
    """ Common implementation for param_get. """
    import lib.commands as commands
    args: dict[str, str | bool | dict[str, str]] = {'uuid': uuid, 'param-name': param_name}
    if key is not None:
        args['param-key'] = key
    try:
        value = host.xe(f'{xe_prefix}-param-get', args)
    except commands.SSHCommandFailed as e:
        if key and accept_unknown_key and e.stdout == "Error: Key %s not found in map" % key:
            value = None
        else:
            raise
    return value

def _param_set(host: Host, xe_prefix: str, uuid: str, param_name: str, value: str | bool | dict[str, str],
               key: str | None = None) -> None:
    """ Common implementation for param_set. """
    args: dict[str, str | bool | dict[str, str]] = {'uuid': uuid}

    if key is not None:
        param_name = '{}:{}'.format(param_name, key)

    args[param_name] = value

    host.xe(f'{xe_prefix}-param-set', args)

def _param_add(host: Host, xe_prefix: str, uuid: str, param_name: str, value: str, key: str | None = None) -> None:
    """ Common implementation for param_add. """
    param_key = f'{key}={value}' if key is not None else value
    args: dict[str, str | bool | dict[str, str]] = {'uuid': uuid, 'param-name': param_name, 'param-key': param_key}

    host.xe(f'{xe_prefix}-param-add', args)

def _param_remove(host: Host, xe_prefix: str, uuid: str, param_name: str, key: str,
                  accept_unknown_key: bool = False) -> None:
    """ Common implementation for param_remove. """
    import lib.commands as commands
    args: dict[str, str | bool | dict[str, str]] = {'uuid': uuid, 'param-name': param_name, 'param-key': key}
    try:
        host.xe(f'{xe_prefix}-param-remove', args)
    except commands.SSHCommandFailed as e:
        if not accept_unknown_key or e.stdout != "Error: Key %s not found in map" % key:
            raise

def _param_clear(host: Host, xe_prefix: str, uuid: str, param_name: str) -> None:
    """ Common implementation for param_clear. """
    args: dict[str, str | bool | dict[str, str]] = {'uuid': uuid, 'param-name': param_name}
    host.xe(f'{xe_prefix}-param-clear', args)

def hash_password(password: str) -> str:
    """Hash password for /etc/shadow."""
    # XCP-ng uses sha512 with 5000 rounds by default
    return sha512_crypt.using(rounds=5000).hash(password)