HOME

sh-3ll 1.0
DIR:/opt/cloudlinux/venv/lib64/python3.11/site-packages/clwizard/
Current File : //opt/cloudlinux/venv/lib64/python3.11/site-packages/clwizard/wizard.py
# coding=utf-8
#
# Copyright © Cloud Linux GmbH & Cloud Linux Software, Inc 2010-2021 All Rights Reserved
#
# Licensed under CLOUD LINUX LICENSE AGREEMENT
# http://cloudlinux.com/docs/LICENCE.TXT
#
import json
import os
import sys
import time
import traceback
from typing import ClassVar, NoReturn

import psutil
from clcommon import FormattedException
from clcommon.utils import (
    ExternalProgramFailed,
    get_cl_version,
    get_package_db_errors,
    is_ubuntu,
    run_command,
)

from clwizard.config import NoSuchModule

from .config import (
    Config,
    acquire_config_access,
)
from .constants import (
    CRASH_LOG_PATH,
    FILE_MARKER_PATH,
    MAIN_LOG_PATH,
    ModuleStatus,
    WizardStatus,
)
from .exceptions import CancelModuleException, InstallationFailedException
from .modules import ALL_MODULES, get_supported_modules, run_installation
from .parser import parse_cloudlinux_wizard_opts
from .utils import (
    is_background_process_running,
    run_background,
    setup_logger,
)


class CloudlinuxWizard:
    """Main class for working with Wizard that exposes high level logic."""

    # states in which we can remove the module from queue
    CANCELLABLE_MODULE_STATUSES: ClassVar[list[str]] = [
        ModuleStatus.PENDING,
        ModuleStatus.FAILED,
        ModuleStatus.CANCELLED,
    ]

    # modules states in which wizard modules can be considered as done
    DONE_MODULES_STATUSES: ClassVar[list[str]] = [
        ModuleStatus.INSTALLED,
        ModuleStatus.CANCELLED,
        ModuleStatus.AUTO_SKIPPED,
    ]

    def __init__(self):
        self._opts = None
        self._supported_modules = get_supported_modules()
        self.log = setup_logger("wizard.main", MAIN_LOG_PATH)

    def run(self, argv):
        """
        CL Wizard main function.

        :param argv: command line arguments for wizard
        :return: None
        """
        self._opts = parse_cloudlinux_wizard_opts(argv)
        try:
            if self._opts.subparser == "install":
                self._validate_system()
                if self.is_installation_finished() and not self._opts.force:
                    self._print_result_and_exit(
                        result="Installation already finished",
                        exit_code=1,
                    )
                if self._opts.no_async:
                    # In async mode, run_background_installation() spawns a new process with --no-async,
                    # which will execute this branch. So we call _prepare_for_installation() here.
                    self._prepare_for_installation()
                    run_installation()
                else:
                    self.run_background_installation(options=self._opts.json_data)
            elif self._opts.subparser == "status":
                self._validate_system()
                if self._opts.initial:
                    self._get_initial_status()
                else:
                    self._get_modules_statuses()
            elif self._opts.subparser == "cancel":
                self._cancel_module_installation(self._opts.module)
            elif self._opts.subparser == "finish":
                self.create_completion_marker()
            else:
                raise NotImplementedError

            if (self._opts.subparser in ["install", "cancel"] and self.is_all_modules_installed()) or (
                self._opts.subparser == "finish" and not self.is_all_modules_installed()
            ):
                # Called only once if:
                #  -- in case of an install:   -all modules were installed successfully
                #                              -a module failed during installation,
                #                              but was installed after resuming
                #  -- in case of cancelling:   -a module failed during installation,
                #                              but was canceled by the user and as a result,
                #                              all modules in a 'done' status
                #  -- in case of finish:       -only if user closed the wizard while a module
                #                              had a status other than installed, cancelled or skipped
                self.run_collecting_statistics()
                self.run_cagefs_force_update()

            self._print_result_and_exit()
        except FormattedException as err:
            self.log.exception(
                "Got an error while running cloudlinux-wizard",
                exc_info=err,
            )
            self._print_result_and_exit(
                result=err.message,
                context=err.context,
                details=err.details,
                exit_code=1,
            )
        except InstallationFailedException:
            self._print_result_and_exit(
                result="Module installation failed, see the log for more information",
                exit_code=1,
            )
        except Exception as err:
            self.log.exception("Unknown error in cloudlinux-wizard", exc_info=err)
            self._print_result_and_exit(
                result="Unknown error occured, please, try again or contact CloudLinux support if it persists.",
                details=traceback.format_exc(),
            )

    @staticmethod
    def is_installation_finished():
        # type: () -> bool
        return os.path.isfile(FILE_MARKER_PATH)

    def create_completion_marker(self):
        # type: () -> None
        try:
            os.mknod(FILE_MARKER_PATH)
            self.log.info("Wizard execution complete")
        except OSError as err:
            self.log.warning(
                "Wizard 'finish' command called more than once, error: '%s'",
                str(err),
            )
            self._print_result_and_exit(
                result="Wizard 'finish' command called more than once",
                exit_code=1,
            )

    def run_background_installation(self, options: dict | None = None) -> None:
        cmd = sys.argv[:]
        cmd.append("--no-async")
        with acquire_config_access() as config:
            # two processes cannot use config at same time
            # so we can safely do check for running process here
            if is_background_process_running():
                self._print_result_and_exit(
                    result="Unable to start a new installation because a background task is still working",
                    exit_code=1,
                )
            # the only case when options are None is the 'resume' case
            if options is not None:
                config.set_modules(options)
            # worker will not be able to acquire reading lock
            # and will wait unless we finally close config file
            worker_pid = run_background(cmd).pid
            config.worker_pid = worker_pid
        self._print_result_and_exit(result="success", pid=worker_pid)

    def _validate_system(self):
        """
        Check whether Wizard supports the current system.
        """
        if get_cl_version() is None:
            self._print_result_and_exit(
                result="Could not identify the CloudLinux version. "
                "Restart your system. If you have the same problem again - "
                "contact CloudLinux support.",
            )

    def _prepare_for_installation(self):
        """
        Prepare the environment before performing the installation.
        This function updates the package lists if run on Ubuntu, expires Yum
        cache when running on RHEL-based systems.
        """
        if is_ubuntu():
            cmd = ["apt-get", "-q", "update"]
            try:
                out = run_command(cmd)
                self.log.info("apt-get update output:\n%s", out)
            except ExternalProgramFailed as err:
                self.log.exception("Error during apt-get update", exc_info=err)
        else:
            cmd = ["yum", "-qy", "clean", "expire-cache"]
            try:
                out = run_command(cmd)
                self.log.info("yum clean expire-cache output:\n%s", out)
            except ExternalProgramFailed as err:
                self.log.exception("Error during yum clean expire-cache", exc_info=err)

    def _get_module_log_path(self, module_name):
        """
        Get path to module log file.
        """
        return self._supported_modules[module_name].LOG_FILE

    def _get_modules_statuses(self):
        """
        Get information about background worker state.
        """
        # we should return modules in order, but config
        # does not know about it, let's sort modules here
        modules = []
        with acquire_config_access() as config:
            state = self._get_wizard_state(config)
            for name in self._supported_modules:
                try:
                    status = config.get_module_status(name)
                    status_time = config.get_module_status_time(name)
                except NoSuchModule:
                    continue
                module_status = {
                    "status": status,
                    "name": name,
                    "status_time": status_time,
                }
                if status in [ModuleStatus.FAILED, ModuleStatus.AUTO_SKIPPED]:
                    module_status["log_file"] = self._get_module_log_path(name)
                modules.append(module_status)
        if state == WizardStatus.CRASHED:
            self._print_result_and_exit(
                wizard_status=state,
                modules=modules,
                crash_log=CRASH_LOG_PATH,
            )
        self._print_result_and_exit(wizard_status=state, modules=modules)

    def _get_initial_status(self):
        """
        Get initial modules status that is used by lvemanager to display wizard pages.
        """
        error_message = get_package_db_errors()
        if error_message:
            # package manager DB corrupted
            self._print_result_and_exit(result=error_message)
        else:
            all_modules_set = {str(key) for key in ALL_MODULES}
            supported_modules_set = set(self._supported_modules.keys())
            unsupported_by_cp = list(all_modules_set - supported_modules_set)
            self._print_result_and_exit(
                modules={module_name: cls().initial_status() for module_name, cls in self._supported_modules.items()},
                unsuppored_by_cp=unsupported_by_cp,
            )

    def _cancel_module_installation(self, module: str) -> None:
        """
        Remove module from queue or print the error if it's not possible.
        """
        self.log.info("Trying to cancel the installation of module '%s'", module)
        with acquire_config_access() as config:
            status = config.get_module_status(module)
            if status in self.CANCELLABLE_MODULE_STATUSES:
                config.set_module_status(
                    module_name=module,
                    new_state=ModuleStatus.CANCELLED,
                )
                self.log.info("Module '%s' installation successfully canceled", module)
            else:
                self.log.warning(
                    "Unable to cancel module '%s' installation, because it is in status '%s'",
                    module,
                    status,
                )
                raise CancelModuleException(module, status)

    def run_collecting_statistics(self):
        """
        Collect user's statistics.
        """
        cmd = ["/usr/sbin/cloudlinux-summary", "--send"]
        if not os.environ.get("SYNCHRONOUS_SUMMARY"):
            cmd.append("--async")

        self.log.info("Collecting statistics...")
        try:
            out = run_command(cmd)
            self.log.info("Statistics collection command output: '%s'", out)
        except ExternalProgramFailed as err:
            self.log.exception("Error during statistics collection", exc_info=err)

    def is_all_modules_installed(self):
        # type: () -> bool
        """
        Check that all modules were either:
        -- installed
        -- canceled
        -- or auto-skipped
        """
        with acquire_config_access() as config:
            statuses = list(config.statuses.values())
        return all(status in self.DONE_MODULES_STATUSES for status in statuses)

    def run_cagefs_force_update(self):
        """
        Run cagefsctl --force-update in background.
        """
        cagefsctl_bin = "/usr/sbin/cagefsctl"

        if not os.path.isfile(cagefsctl_bin):
            return
        cmd = [cagefsctl_bin, "--force-update", "--wait-lock"]
        self.log.info("Starting cagefs force-update in the background: %s", cmd)
        cagefsctl_proc = run_background(cmd)
        # In Cloudlinux tests environment statistics wait for cagefsctl --force-update terminate
        is_test_environment = bool(os.environ.get("CL_TEST_SYSTEM"))
        if is_test_environment:
            cagefsctl_proc.communicate()

    def _get_wizard_state(self, config: Config) -> str:
        # worker pid is None only in the case when wizard
        # wasn't EVER called, this worker pid will stay
        # in config forever, even after wizard is Done
        if config.worker_pid is None:
            return WizardStatus.IDLE

        try:
            psutil.Process(config.worker_pid)
        except psutil.NoSuchProcess:
            # Background process is no longer alive.
            # 1. Wizard DONE: all modules are in state "installed", "cancelled" or "auto-skipped".
            # 2. Wizard FAILED: one of the modules in state "failed" or "cancelled"
            # and no modules are in status "installing"
            # 3. Wizard CRASHED: none of the above.
            statuses = list(config.statuses.values())
            if all(status in self.DONE_MODULES_STATUSES for status in statuses):
                return WizardStatus.DONE
            # cancel module`s status is acceptable for general wizard status FAILED, DO NOT CHANGE IT PLS (LU-1295)
            # An extra check for "installing" status is needed to exclude possible CRASHED wizard status
            if any(status in (ModuleStatus.FAILED, ModuleStatus.CANCELLED) for status in statuses) and not any(
                status in (ModuleStatus.INSTALLING,) for status in statuses
            ):
                return WizardStatus.FAILED
            return WizardStatus.CRASHED

        return WizardStatus.IN_PROGRESS

    @staticmethod
    def _print_result_and_exit(
        result: str = "success",
        exit_code: int = 0,
        **extra,
    ) -> NoReturn:
        """
        Print data in default format for web and exit.

        :param dict extra: extra fields for the response, usually we expect 'context' here
        """
        message = {"result": result, "timestamp": time.time()}
        message.update(extra)
        print(json.dumps(message, indent=2, sort_keys=True))
        sys.exit(exit_code)