1#@!#!123s

: / lib / python3 / dist-packages / uaclient / cli /
Filename : cli_util.py
back
import itertools
import re
import sys
import threading
import time
from functools import wraps
from typing import Optional

from uaclient import (
    actions,
    api,
    daemon,
    entitlements,
    event_logger,
    exceptions,
    lock,
    messages,
    status,
    util,
)
from uaclient.api.u.pro.security.fix._common import CVE_OR_USN_REGEX
from uaclient.api.u.pro.status.is_attached.v1 import _is_attached
from uaclient.apt import AptProxyScope, setup_apt_proxy
from uaclient.config import UAConfig
from uaclient.files import machine_token

event = event_logger.get_event_logger()


CLEAR_LINE_ANSI_CODE = "\033[K"


class CLIEnableDisableProgress(api.AbstractProgress):
    def __init__(self, *, assume_yes: bool):
        self.is_interactive = not assume_yes
        self.assume_yes = assume_yes

    def progress(
        self,
        *,
        total_steps: int,
        done_steps: int,
        previous_step_message: Optional[str],
        current_step_message: Optional[str]
    ):
        if current_step_message is not None:
            print(current_step_message)

    def _on_event(self, event: str, payload):
        if event == "info":
            print(payload)
        elif event == "message_operation":
            if not util.handle_message_operations(payload, self.assume_yes):
                raise exceptions.PromptDeniedError()


def _null_print(*args, **kwargs):
    pass


def create_interactive_only_print_function(json_output: bool):
    if json_output:
        return _null_print
    else:
        return print


def assert_lock_file(lock_holder=None):
    """Decorator asserting exclusive access to lock file"""

    def wrapper(f):
        @wraps(f)
        def new_f(*args, cfg, **kwargs):
            with lock.RetryLock(lock_holder=lock_holder, sleep_time=1):
                retval = f(*args, cfg=cfg, **kwargs)
            return retval

        return new_f

    return wrapper


def assert_root(f):
    """Decorator asserting root user"""

    @wraps(f)
    def new_f(*args, **kwargs):
        if not util.we_are_currently_root():
            raise exceptions.NonRootUserError()
        else:
            return f(*args, **kwargs)

    return new_f


def verify_json_format_args(f):
    """Decorator to verify if correct params are used for json format"""

    @wraps(f)
    def new_f(cmd_args, *args, **kwargs):
        if not cmd_args:
            return f(cmd_args, *args, **kwargs)

        if cmd_args.format == "json" and not cmd_args.assume_yes:
            raise exceptions.CLIJSONFormatRequireAssumeYes()
        else:
            return f(cmd_args, *args, **kwargs)

    return new_f


def assert_attached(raise_custom_error_function=None):
    """Decorator asserting attached config.
    :param msg_function: Optional function to generate a custom message
    if raising an UnattachedError
    """

    def wrapper(f):
        @wraps(f)
        def new_f(args, cfg, **kwargs):
            if not _is_attached(cfg).is_attached:
                if raise_custom_error_function:
                    command = getattr(args, "command", "")
                    service_names = getattr(args, "service", "")
                    raise_custom_error_function(
                        command=command, service_names=service_names, cfg=cfg
                    )
                else:
                    raise exceptions.UnattachedError()
            return f(args, cfg=cfg, **kwargs)

        return new_f

    return wrapper


def assert_not_attached(f):
    """Decorator asserting unattached config."""

    @wraps(f)
    def new_f(args, cfg, **kwargs):
        if _is_attached(cfg).is_attached:
            machine_token_file = machine_token.get_machine_token_file()
            raise exceptions.AlreadyAttachedError(
                account_name=machine_token_file.account.get("name", "")
            )
        return f(args, cfg=cfg, **kwargs)

    return new_f


def assert_vulnerability_issue_valid(cmd):
    def wrapper(f):
        @wraps(f)
        def new_f(args, cfg, **kwargs):
            security_issue = getattr(args, "security_issue", "")
            if not re.match(CVE_OR_USN_REGEX, security_issue):
                raise exceptions.InvalidSecurityIssueIdFormat(
                    issue=security_issue,
                    cmd=cmd,
                )

            return f(args, cfg=cfg, **kwargs)

        return new_f

    return wrapper


def _raise_enable_disable_unattached_error(command, service_names, cfg):
    """Raises a custom error for enable/disable commands when unattached.

    Takes into consideration if the services exist or not, and notify the user
    accordingly."""
    (
        entitlements_found,
        entitlements_not_found,
    ) = entitlements.get_valid_entitlement_names(names=service_names, cfg=cfg)
    if entitlements_found and entitlements_not_found:
        raise exceptions.UnattachedMixedServicesError(
            valid_service=", ".join(entitlements_found),
            operation=command,
            invalid_service=", ".join(entitlements_not_found),
            service_msg="",
        )
    elif entitlements_found:
        raise exceptions.UnattachedValidServicesError(
            valid_service=", ".join(entitlements_found), operation=command
        )
    else:
        raise exceptions.UnattachedInvalidServicesError(
            operation=command,
            invalid_service=", ".join(entitlements_not_found),
            service_msg="",
        )


def post_cli_attach(cfg: UAConfig) -> None:
    machine_token_file = machine_token.get_machine_token_file(cfg)
    contract_name = machine_token_file.contract_name

    if contract_name:
        event.info(
            messages.ATTACH_SUCCESS_TMPL.format(contract_name=contract_name)
        )
    else:
        event.info(messages.ATTACH_SUCCESS_NO_CONTRACT_NAME)

    daemon.stop()
    daemon.cleanup(cfg)

    status_dict, _ret = actions.status(cfg)
    output = status.format_tabular(status_dict)
    event.info(util.handle_unicode_characters(output))
    event.process_events()


def configure_apt_proxy(
    cfg: UAConfig,
    scope: AptProxyScope,
    set_key: str,
    set_value: Optional[str],
) -> None:
    """
    Handles setting part the apt proxies - global and uaclient scoped proxies
    """
    if scope == AptProxyScope.GLOBAL:
        http_proxy = cfg.global_apt_http_proxy
        https_proxy = cfg.global_apt_https_proxy
    elif scope == AptProxyScope.UACLIENT:
        http_proxy = cfg.ua_apt_http_proxy
        https_proxy = cfg.ua_apt_https_proxy
    if "https" in set_key:
        https_proxy = set_value
    else:
        http_proxy = set_value
    setup_apt_proxy(
        http_proxy=http_proxy, https_proxy=https_proxy, proxy_scope=scope
    )


def run_spinner(stop_spinner, msg):
    spinner = itertools.cycle(["|", "/", "-", "\\"])
    while not stop_spinner.is_set():
        sys.stdout.write(
            "\r{}{}{}".format(msg, next(spinner), CLEAR_LINE_ANSI_CODE)
        )
        sys.stdout.flush()
        time.sleep(0.1)
        sys.stdout.write("\b")


def with_spinner(msg: Optional[str] = ""):
    def wrapper(f):
        @wraps(f)
        def new_f(args, *, cfg, **kwargs):
            if not sys.stdout.isatty():
                return f(args, cfg=cfg, **kwargs)

            stop_spinner = threading.Event()

            spinner_thread = threading.Thread(
                target=run_spinner, args=(stop_spinner, msg)
            )
            spinner_thread.start()

            retval = f(args, cfg=cfg, **kwargs)

            stop_spinner.set()
            spinner_thread.join()

            sys.stdout.write("\r" + CLEAR_LINE_ANSI_CODE)
            sys.stdout.flush()

            return retval

        return new_f

    return wrapper


def colorize_priority(priority):
    if priority == "low":
        return messages.TxtColor.INFOBLUE + priority + messages.TxtColor.ENDC
    elif priority == "medium":
        return (
            messages.TxtColor.WARNINGYELLOW + priority + messages.TxtColor.ENDC
        )
    elif priority == "high":
        return messages.TxtColor.ORANGE + priority + messages.TxtColor.ENDC
    elif priority == "critical":
        return messages.TxtColor.FAIL + priority + messages.TxtColor.ENDC
    else:
        return priority