Source code for watchmaker.managers.platform_manager

# -*- coding: utf-8 -*-
"""Watchmaker base manager."""
from __future__ import (
    absolute_import,
    division,
    print_function,
    unicode_literals,
    with_statement,
)

import concurrent.futures
import logging
import os
import shutil
import subprocess
import tarfile
import tempfile
import zipfile

import watchmaker.utils
from watchmaker.exceptions import WatchmakerError
from watchmaker.utils import urllib_utils




class PlatformManagerBase(object):
    """
    Base class for operating system managers.

    All child classes will have access to methods unless overridden by an
    identically-named method in the child class.

    Args:
        system_params: (:obj:`dict`)
            Attributes, mostly file-paths, specific to the system-type (Linux
            or Windows). The dict keys are as follows:

            prepdir:
                Directory where Watchmaker will keep files on the system.
            readyfile:
                Path to a file that will be created upon successful
                completion.
            logdir:
                Directory to store log files.
            workingdir:
                Directory to store temporary files. Deleted upon successful
                completion.
            restart:
                Command to use to restart the system upon successful
                completion.
            shutdown_path:
                (Windows-only) Path to the Windows ``shutdown.exe`` command.

    """

    boto3 = None
    boto_client = None

    def __init__(self, system_params, *args, **kwargs):
        self.log = logging.getLogger(
            '{0}.{1}'.format(__name__, self.__class__.__name__)
        )
        self.system_params = system_params
        self.working_dir = None
        PlatformManagerBase.args = args
        PlatformManagerBase.kwargs = kwargs



    def retrieve_file(self, url, filename):
        """
        Retrieve a file from a provided URL.

        Supports all :obj:`urllib.request` handlers, as well as S3 buckets.

        Args:
            url: (:obj:`str`)
                URL to a file.

            filename: (:obj:`str`)
                Path where the file will be saved.

        """
        # Convert a local path to a URI
        url = watchmaker.utils.uri_from_filepath(url)
        self.log.debug('Downloading: %s', url)
        self.log.debug('Destination: %s', filename)

        try:
            self.log.debug('Establishing connection to the host, %s', url)
            response = watchmaker.utils.urlopen_retry(url)
            self.log.debug('Opening the file handle, %s', filename)
            with open(filename, 'wb') as outfile:
                self.log.debug('Saving file to local filesystem...')
                shutil.copyfileobj(response, outfile)
        except (ValueError, urllib_utils.error.URLError):
            self.log.critical(
                'Failed to retrieve the file. url = %s. filename = %s',
                url, filename
            )
            raise
        self.log.info(
            'Retrieved the file successfully. url=%s. filename=%s',
            url, filename
        )




    def create_working_dir(self, basedir, prefix):
        """
        Create a directory in ``basedir`` with a prefix of ``prefix``.

        Args:
            prefix: (:obj:`str`)
                Prefix to prepend to the working directory.

            basedir: (:obj:`str`)
                The directory in which to create the working directory.

        Returns:
            :obj:`str`: Path to the working directory.

        """
        self.log.info('Creating a working directory.')
        original_umask = os.umask(0)
        try:
            working_dir = tempfile.mkdtemp(prefix=prefix, dir=basedir)
        except Exception:
            msg = 'Could not create a working dir in {0}'.format(basedir)
            self.log.critical(msg)
            raise
        self.log.debug('Created working directory: %s', working_dir)
        os.umask(original_umask)
        return working_dir


    @staticmethod
    def _pipe_handler(pipe, logger=None, prefix_msg=''):
        ret = b''
        try:
            for line in iter(pipe.readline, b''):
                if logger:
                    logger('%s%s', prefix_msg, line.rstrip())
                ret += line
        finally:
            pipe.close()

        return ret



    def call_process(self, cmd, log_pipe='all', raise_error=True):
        """
        Execute a shell command.

        Args:
            cmd: (:obj:`list`)
                Command to execute.

            log_pipe: (:obj:`str`)
                Controls what to log from the command output. Supports three
                values: ``stdout``, ``stderr``, ``all``.
                (*Default*: ``all``)

            raise_error: (:obj:`bool`)
                Switch to control whether to raise if the command return code
                is non-zero.
                (*Default*: ``True``)

        Returns:
            :obj:`dict`:
                Dictionary containing three keys: ``retcode`` (:obj:`int`),
                ``stdout`` (:obj:`bytes`), and ``stderr`` (:obj:`bytes`).

        """
        ret = {
            'retcode': 0,
            'stdout': b'',
            'stderr': b''
        }

        if not isinstance(cmd, list):
            msg = 'Command is not a list: {0}'.format(cmd)
            self.log.critical(msg)
            raise WatchmakerError(msg)

        self.log.debug('Command: %s', ' '.join(cmd))

        # If running as a standalone, PyInstaller will have modified the
        # LD_LIBRARY_PATH to point to standalone libraries. If there were a
        # value at runtime, PyInstaller will create LD_LIBRARY_PATH_ORIG. In
        # order for salt to run correctly, LD_LIBRARY_PATH has to be fixed.
        kwargs = {}
        env = dict(os.environ)
        lib_path_key = 'LD_LIBRARY_PATH'

        if env.get(lib_path_key) is not None:

            lib_path_orig_value = env.get(lib_path_key + '_ORIG')
            if lib_path_orig_value is None:

                # you can have lib_path and no orig, if:
                # 1. none was set and pyinstaller set one, or
                # 2. one was set and we're not in standalone package
                env.pop(lib_path_key, None)

            else:

                # put original lib_path back
                env[lib_path_key] = lib_path_orig_value

            kwargs['env'] = env

        with subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kwargs
        ) as process, concurrent.futures.ThreadPoolExecutor(
            max_workers=2
        ) as executor:
            stdout_future = executor.submit(
                self._pipe_handler,
                process.stdout,
                self.log.debug if log_pipe in ['stdout', 'all'] else None,
                'Command stdout: '
            )

            stderr_future = executor.submit(
                self._pipe_handler,
                process.stderr,
                self.log.error if log_pipe in ['stderr', 'all'] else None,
                'Command stderr: ')

            ret['stdout'] = stdout_future.result()
            ret['stderr'] = stderr_future.result()

            ret['retcode'] = process.wait()

        self.log.debug('Command retcode: %s', ret['retcode'])

        if raise_error and ret['retcode'] != 0:
            msg = 'Command failed! Exit code={0}, cmd={1}'.format(
                ret['retcode'], ' '.join(cmd))
            self.log.critical(msg)
            raise WatchmakerError(msg)

        return ret




    def cleanup(self):
        """Delete working directory."""
        self.log.info('Cleanup Time...')
        try:
            self.log.debug('working_dir=%s', self.working_dir)
            shutil.rmtree(self.working_dir)
            self.log.info('Deleted working directory...')
        except Exception:
            msg = 'Cleanup Failed!'
            self.log.critical(msg)
            raise

        self.log.info('Exiting cleanup routine...')




    def extract_contents(self, filepath, to_directory, create_dir=False):
        """
        Extract a compressed archive to the specified directory.

        Args:
            filepath: (:obj:`str`)
                Path to the compressed file. Supported file extensions:

                - `.zip`
                - `.tar.gz`
                - `.tgz`
                - `.tar.bz2`
                - `.tbz`

            to_directory: (:obj:`str`)
                Path to the target directory

            create_dir: (:obj:`bool`)
                Switch to control the creation of a subdirectory within
                ``to_directory`` named for the filename of the compressed file.
                (*Default*: ``False``)

        """
        if filepath.endswith('.zip'):
            self.log.debug('File Type: zip')
            opener, mode = zipfile.ZipFile, 'r'
        elif filepath.endswith('.tar.gz') or filepath.endswith('.tgz'):
            self.log.debug('File Type: GZip Tar')
            opener, mode = tarfile.open, 'r:gz'
        elif filepath.endswith('.tar.bz2') or filepath.endswith('.tbz'):
            self.log.debug('File Type: Bzip Tar')
            opener, mode = tarfile.open, 'r:bz2'
        else:
            msg = (
                'Could not extract "{0}" as no appropriate extractor is found.'
                .format(filepath)
            )
            self.log.critical(msg)
            raise WatchmakerError(msg)

        if create_dir:
            to_directory = os.sep.join((
                to_directory,
                '.'.join(filepath.split(os.sep)[-1].split('.')[:-1])
            ))

        try:
            os.makedirs(to_directory)
        except OSError:
            if not os.path.isdir(to_directory):
                msg = 'Unable create directory - {0}'.format(to_directory)
                self.log.critical(msg)
                raise

        cwd = os.getcwd()
        os.chdir(to_directory)

        try:
            openfile = opener(filepath, mode)
            try:
                openfile.extractall()
            finally:
                openfile.close()
        finally:
            os.chdir(cwd)

        self.log.info(
            'Extracted file. source=%s, dest=%s',
            filepath, to_directory
        )






class LinuxPlatformManager(PlatformManagerBase):
    """
    Base class for Linux Platforms.

    Serves as a foundational class to keep OS consistency.
    """

    def _install_from_yum(self, packages):
        yum_cmd = ['sudo', 'yum', '-y', 'install']
        if isinstance(packages, list):
            yum_cmd.extend(packages)
        else:
            yum_cmd.append(packages)
        self.call_process(yum_cmd)
        self.log.debug(packages)





class WindowsPlatformManager(PlatformManagerBase):
    """
    Base class for Windows Platform.

    Serves as a foundational class to keep OS consistency.
    """