Mini Shell

import asyncio
import logging
import functools
import os

from defence360agent.subsys.panels.base import PanelException
from defence360agent.subsys.panels.directadmin.config import ConfigOptions
from defence360agent.utils import check_run, CheckRunError

logger = logging.getLogger(__name__)


BASEDIR = "/usr/local/directadmin/custombuild"
BUILD = os.path.join(BASEDIR, "build")


class CustomBuildOptions(ConfigOptions):
    DEFAULT_FILENAME = os.path.join(BASEDIR, "options.conf")


async def build(option):
    MAX_RETRIES = 3
    for attempt in range(MAX_RETRIES + 1):
        try:
            return await check_run([BUILD, option], cwd=BASEDIR)
        except CheckRunError as e:
            if attempt == MAX_RETRIES:
                raise PanelException(
                    "Failed to run custombuild command '{} {}' after {}"
                    " attempts: {}".format(BUILD, option, MAX_RETRIES, e)
                )
            if "Another instance of custombuild" in str(e):
                logger.warning(
                    "Failed to run custombuild command '%s %s': %s",
                    BUILD,
                    option,
                    e,
                )
                logger.info("Retrying custombuild command after delay...")
                await asyncio.sleep(5)
            else:
                raise PanelException(
                    "Failed to run custombuild command '{} {}': {}".format(
                        BUILD, option, e
                    )
                )


def custombuild2_only(fun):
    @functools.wraps(fun)
    async def wrapper(*args, **kwargs):
        try:
            out = await build("version")
        except (CheckRunError, FileNotFoundError, PermissionError):
            pass
        else:
            if out.startswith(b"2."):
                return await fun(*args, **kwargs)
        raise PanelException(
            "{} is only supported on installations "
            "with Custombuild 2.0".format(fun.__name__)
        )

    return wrapper