OwlCyberSecurity - MANAGER

Edit File: progress_check.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/LICENSE.TXT

import functools
import json
from pathlib import Path
from typing import Dict

from clwpos.cl_wpos_exceptions import WposError
from clwpos.constants import USER_WPOS_DIR
from clwpos.logsetup import setup_logging
from clwpos.optimization_features import ALL_OPTIMIZATION_FEATURES
from clwpos.utils import home_dir, is_run_under_user
from secureio import write_file_via_tempfile

logger = setup_logging(__name__)

PROGRESS_STATUS_FILE_NAME = '.progress-status'

def get_progress_file_path() -> Path:
    return Path(home_dir(), USER_WPOS_DIR, PROGRESS_STATUS_FILE_NAME)

class CommandProgress:
    """
    Class to represent user's AccelerateWP utility commands progress.
    """
    total_stages: int

def __new__(cls, command):
        command_subclass_map = {
            subclass.command: subclass for subclass in cls.__subclasses__()
        }
        if command not in command_subclass_map:
            raise WposError(f'Internal Error: command "{command}" is not allowed, '
                            f'available commands: {list(command_subclass_map.keys())}')
        subclass = command_subclass_map[command]
        instance = super(CommandProgress, subclass).__new__(subclass)
        return instance

def __init__(self, *args):
        if not is_run_under_user():
            raise WposError('Internal Error: trying to use CommandProgress class as root')
        self.progress_file = get_progress_file_path()
        self.in_progress = False
        self.completed_stages = 0

def as_dict(self) -> Dict[str, int]:
        """
        Return number of total and completed stages as dict.
        """
        return {
            'total_stages': self.total_stages,
            'completed_stages': self.completed_stages,
        }

def start(self) -> None:
        """
        Start tracking the progress of the command.
        """
        self.in_progress = True
        self.completed_stages = 0
        self.progress_file.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
        self._update_progress_file()

def update(self) -> None:
        """
        Update progress status by incrementing
        the number of completed stages by one.
        """
        if not self.in_progress:
            return
        self.completed_stages += 1
        if self.completed_stages > self.total_stages:
            raise WposError(
                'Internal Error: the number of completed stages '
                'is greater than the total number of stages'
            )
        self._update_progress_file()

def complete(self) -> None:
        """
        Complete progress, delete file with progress status.
        """
        if not self.in_progress:
            return
        self.in_progress = False
        self.completed_stages = self.total_stages
        if self.progress_file.exists():
            self.progress_file.unlink()

def _update_progress_file(self) -> None:
        # TODO: improve me
        # clcaptain cannot write via tmp file, disable_quota() also does not work as this code is run only under user
        # (get/enable/disable) not that critical now, because feature installing will not work anyway for overquota user
        write_file_via_tempfile(
            json.dumps(self.as_dict()), self.progress_file, 0o644)

@staticmethod
    def get_status() -> Dict[str, int]:
        """
        Return current progress status of the command execution,
        getting it from from special file.
        """
        progress_file = get_progress_file_path()
        defaults = {'total_stages': 1, 'completed_stages': 0}

if not progress_file.exists():
            return defaults

try:
            with open(progress_file, 'r') as f:
                status_from_file = json.loads(f.read())
            return {
                'total_stages': int(status_from_file['total_stages']),
                'completed_stages': int(status_from_file['completed_stages']),
            }
        except Exception as e:
            logger.error('Can\'t parse progress status json: %s', e)
            return defaults

class EnableCommandProgress(CommandProgress):
    """
    Class to represent user's AccelerateWP utility "enable" command progress.
    """
    command: str = 'enable'
    total_stages: int = 9

class DisableCommandProgress(CommandProgress):
    """
    Class to represent user's WPOS utility "disable" command progress.
    """
    command: str = 'disable'
    total_stages: int = 5

class GetCommandProgress(CommandProgress):
    """
    Class to represent user's WPOS utility "get" command progress.
    """
    command: str = 'get'
    total_stages: int = 10

def recalculate_number_of_total_stages(self, user_info: dict):
        """
        Recalculate and change total number of progress stages depending on
        the number of user's docroots and wordpresses.
        so we could update progress after checking every docroot and wordpress.
        Example of user_info structure:
            {
                'docroot1':
                    {'wps': [{'path': '', 'version': '5.9'},
                             {'path': 'wp2', 'version': '5.9'}], ... },
                'docroot2':
                    {'wps': [{'path': '', 'version': '5.9'}], ... },
            }
        """
        number_of_initial_stages = 3
        number_of_docroots = len(user_info)
        number_of_wordpress = sum(
            len(doc_root_info['wps']) for doc_root_info in user_info.values()
        )
        total_checks_per_module = number_of_docroots + number_of_wordpress
        self.total_stages = number_of_initial_stages \
            + total_checks_per_module * len(ALL_OPTIMIZATION_FEATURES)

def track_command_progress(func):
    """
    Decorator for the WPOS utility commands
    we want to track the progress for.
    """
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        self.command_progress = CommandProgress(func.__name__.strip('_'))
        self.command_progress.start()
        try:
            return func(self, *args, **kwargs)
        finally:
            self.command_progress.complete()
    return wrapper

def update_command_progress(func):
    """
    Decorator for the WPOS utility methods after which
    the counter of completed operations must be updated.
    """
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        result = func(self, *args, **kwargs)
        if getattr(self, 'command_progress') is not None:
            self.command_progress.update()
        return result
    return wrapper