# Copyright 2019 VMware, Inc.  All rights reserved. -- VMware Confidential
#
'''
This module provides the local transport functionality needed to perform the base
operations on target machine.
'''
import sys
import platform
import logging
import os
import glob
import subprocess
import shutil
import time
import uuid
import tempfile
import errno
from zipfile import ZipFile

from status_reporting_sdk.filelock import SecureOpen
import transport
from transport import OperationsManager, ProcessResult, ExecutionException, FileException, ErrorCode
from transport.utils import ensureFileCanBeCreated, ensureFileExists
from shutdown_executor import shutdownHookExecutor
import shutdown_executor

__author__ = "VMware, Inc."

logger = logging.getLogger(__name__)

PY2 = sys.version_info[0] == 2

def killProcess(pid, sig=transport.CTRL_BREAK_SIGNAL):
    try:
        logger.info('Sending signal %s to process %s', sig, pid)
        transport.kill(pid, sig)
        logger.info('Signal %s has been successfully sent to process %s', sig, pid)
    except OSError as e:
        logger.warning("Unable to send signal %s to process %s. Error: %s",
                       sig, pid, str(e))

class _Process(transport.SerializablePopenMixin):
    '''Extend Process result with local popen handle.
    '''
    def __init__(self, popenHandle, startedAt, stdoutTempFile=None, stderrTempFile=None):
        self.processUuid = str(uuid.uuid4())
        self._popenHandle = popenHandle
        self.startedAt = startedAt
        self.stdoutTempFile = stdoutTempFile
        self.stderrTempFile = stderrTempFile

    def getPid(self):
        '''Gets started process id.
        '''
        pid = self._popenHandle.pid
        return pid

    def wait(self):
        '''Waits until the process complete.
        '''
        self._popenHandle.wait()

    def isCompleted(self):
        '''Checks if the process has completed.
        '''
        return self._popenHandle.poll() is not None

    def getProcessResult(self, waitForCompletion=True, keepAsBytes=PY2):
        pid = self.getPid()
        exitCode = None
        # Wait the process to finish if needed
        if waitForCompletion:
            exitCode = self._popenHandle.wait()
        else:
            exitCode = self._popenHandle.poll()
        completed = exitCode is not None
        # Return ouput and error as part of the result only when they haven't
        # been passed, otherwise they can be read from the passed files
        output = None
        err = None
        readMode = 'rb' if keepAsBytes else 'r'
        if self.stdoutTempFile:
            with open(self.stdoutTempFile, readMode) as f:
                output = f.read()
            if completed:
                os.remove(self.stdoutTempFile)
        if self.stderrTempFile:
            with open(self.stderrTempFile, readMode) as f:
                err = f.read()
            if completed:
                os.remove(self.stderrTempFile)

        # Finish time may not be quite accurate.
        return ProcessResult(
            pid,
            startTime=self.startedAt,
            exitCode=exitCode,
            finishTime=time.time() if completed else None,
            stdout=output,
            stderr=err)

class LocalOperationsManager(OperationsManager):
    '''Local operations manager implementation.
    '''
    # Lookup of Proceess Uuid to _LocalProcessResult
    allProcesses = {}

    def __init__(self):
        # Call super constructor
        super(LocalOperationsManager, self).__init__('localhost')

    def _isWindowsPlatform(self):
        '''Check if we are running on Windows.
        '''
        return self.getPlatform().lower() == 'windows'

    def _samefile(self, src, dst):
        '''Check if this is same file path.
        '''
        # Macintosh, Unix.
        if hasattr(os.path, 'samefile'):
            try:
                return os.path.samefile(src, dst)
            except OSError:
                pass

        # Check for same pathname.
        return (os.path.normcase(os.path.abspath(src)) ==
                os.path.normcase(os.path.abspath(dst)))

    def createTempFile(self, prefix='tmp', suffix=''):
        '''
        @see: transport.createTempFile
        '''
        with tempfile.NamedTemporaryFile(suffix=suffix, prefix=prefix, delete=False) as tmp:
            absPath = os.path.abspath(tmp.name)
            return absPath

    def createTempDirectory(self, prefix='tmp', suffix = ''):
        '''
        @see: transport.createTempDirectory
        '''
        return tempfile.mkdtemp(suffix=suffix, prefix=prefix)

    def downloadFileContent(self, targetPath, keepAsBytes=PY2):
        '''
        @see: transport.downloadFileContent
        '''
        ensureFileExists(targetPath, isFileExpected=True, isLocal=False)

        readMode = 'rb' if keepAsBytes else 'r'
        with open(targetPath, readMode) as fp:
            return fp.read()

    def downloadFile(self, targetPath, localPath, fileMode=None,
                     createLocalPath=True):
        '''
        @see: transport.downloadFile
        '''
        if self._samefile(localPath, targetPath):
            logger.info('Skip downloading file to same path %s', localPath)
            return

        # Check input arguments
        ensureFileExists(targetPath, isFileExpected=True, isLocal=False)
        ensureFileCanBeCreated(localPath, createLocalPath,
                                     isFileExpected=True, isLocal=True)

        try:
            # Copy the file
            with open(targetPath, "rb") as fsrc:
                with SecureOpen(localPath, open, True, "wb") as fdst:
                    shutil.copyfileobj(fsrc, fdst)

            # Set the file mode if specified
            if fileMode is not None:
                logger.debug('Set mode %s to file %s', fileMode, localPath)
                os.chmod(localPath, fileMode)
        except OSError as e:
            error = 'Cannot copy file from %s to %s. Error: %s' %\
                                (targetPath, localPath, str(e))
            logger.error(error)
            raise FileException(error)

    def createDirectory(self, targetPath):
        '''
        @see: transport.createDirectory
        '''
        try:
            os.makedirs(targetPath)
        except OSError as e:
            if e.errno == errno.EEXIST and os.path.isdir(targetPath):
                # Skip exception if already created
                pass
            else:
                error = "Cannot create directory %s. Error: %s" %\
                                                (targetPath, str(e))
                logger.error(error)
                raise FileException(error)

    def downloadDirectory(self, targetPath, localPath, createLocalPath=True):
        '''
        @see: transport.downloadDirectory
        '''
        if self._samefile(localPath, targetPath):
            logger.info('Skip downloading directory to same path %s', localPath)
            return

        # Check input arguments
        ensureFileExists(targetPath, isFileExpected=False, isLocal=False)
        ensureFileCanBeCreated(localPath, createLocalPath,
                                     isFileExpected=False, isLocal=True)

        try:
            # Copy the directory
            shutil.copytree(targetPath, localPath)
        except OSError as e:
            error = 'Cannot copy directory from %s to %s. Error: %s' %\
                                (targetPath, localPath, str(e))
            logger.error(error)
            raise FileException(error)

    def uploadFileContent(self, fileContent, targetPath, fileMode=None,
                          createTargetPath=True):
        '''
        @see: transport.uploadFileContent
        '''
        ensureFileCanBeCreated(targetPath, createTargetPath,
                                     isFileExpected=True, isLocal=False)

        if self.isFile(targetPath):
            logger.info('Content of file %s will be overwritten', targetPath)

        try:
            # Set the file content
            with SecureOpen(targetPath, open, True, 'wb') as fp:
                fp.write(fileContent)

            # Set the file mode if specified
            if fileMode is not None:
                logger.debug('Set mode %s to file %s', fileMode, targetPath)
                os.chmod(targetPath, fileMode)
        except OSError as e:
            error = 'Cannot copy file content to %s. Error: %s' %\
                                (targetPath, str(e))
            logger.error(error)
            raise FileException(error)

    def uploadFile(self, localPath, targetPath, fileMode=None,
                   createTargetPath=True):
        '''
        @see: transport.uploadFile
        '''
        if self._samefile(localPath, targetPath):
            logger.info('Skip uploading file to same path %s', localPath)
            return

        # Check input arguments
        ensureFileExists(localPath, isFileExpected=True, isLocal=True)
        ensureFileCanBeCreated(targetPath, createTargetPath,
                                     isFileExpected=True, isLocal=False)

        try:
            # Copy file
            with open(localPath, "rb") as fsrc:
                with SecureOpen(targetPath, open, True, "wb") as fdst:
                    shutil.copyfileobj(fsrc, fdst)

            # Change the file mode
            if fileMode is not None:
                logger.debug('Set mode %s to file %s', fileMode, targetPath)
                os.chmod(targetPath, fileMode)
        except OSError as e:
            error = 'Cannot copy file from %s to %s. Error: %s' %\
                                (localPath, targetPath, str(e))
            logger.error(error)
            raise FileException(error)

    def uploadDirectory(self, localPath, targetPath, createTargetPath=True):
        '''
        @see: transport.uploadDirectory
        '''
        if self._samefile(localPath, targetPath):
            logger.info("Skip uploading directory to same path '%s'", localPath)
            return

        # Check input arguments
        ensureFileExists(localPath, isFileExpected=False, isLocal=True)
        ensureFileCanBeCreated(targetPath, createTargetPath,
                                     isFileExpected=False, isLocal=False)

        try:
            # Copy the directory
            shutil.copytree(localPath, targetPath)
        except OSError as e:
            error = 'Cannot copy directory from %s to %s. Error: %s' %\
                            (localPath, targetPath, str(e))
            logger.error(error)
            raise FileException(error)

    def pathExists(self, targetPath):
        '''
        @see: transport.fileExists
        '''
        return os.path.exists(targetPath)

    def isFile(self, targetPath):
        '''
        @see: transport.isFile
        '''
        return os.path.isfile(targetPath)

    def isDirectory(self, targetPath):
        '''
        @see: transport.isDirectory
        '''
        return os.path.isdir(targetPath)

    def removeFile(self, targetPath):
        '''
        @see: transport.removeFile
        '''
        if not self.pathExists(targetPath):
            logger.info("File %s does not exist. Skip removing it", targetPath)
        else:
            ensureFileExists(targetPath, isFileExpected=True, isLocal=False)
            try:
                # Remove the file
                os.remove(targetPath)
            except OSError as e:
                error = 'Cannot remove file %s. Error: %s' %\
                                    (targetPath, str(e))
                logger.error(error)
                raise FileException(error)

    def removeDirectory(self, targetPath):
        '''
        @see: transport.removeDirectory
        '''
        if not self.pathExists(targetPath):
            logger.info("Directory %s does not exist. Skip removing it", targetPath)
        else:
            ensureFileExists(targetPath, isFileExpected=False, isLocal=False)
            try:
                # Remove the directory
                shutil.rmtree(targetPath)
            except OSError as e:
                error = 'Cannot remove directory %s. Error: %s' %\
                                    (targetPath, str(e))
                logger.error(error)
                raise FileException(error)

    def listFiles(self, targetPath):
        '''
        @see: transport.listFiles
        '''
        regexPath = "%s%s*" % (targetPath, os.sep)
        result = glob.glob(regexPath)
        return sorted(result)

    def zipFile(self, targetPath, localZipFilePath, createLocalPath=True):
        '''
        @see: transport.zipFile
        '''
        if not self.isFile(targetPath):
            error = 'Filepath %s is not a regular file' % targetPath
            logger.warning(error)
            raise FileException(error)

        ensureFileCanBeCreated(localZipFilePath, createLocalPath,
                                     isFileExpected=True, isLocal=True)

        zipWriter = None
        try:
            zipWriter = ZipFile(localZipFilePath, 'w')
            logger.debug("Handle zipping a file %s to archive %s",
                         targetPath, localZipFilePath)
            zipWriter.write(targetPath)
        finally:
            if zipWriter:
                zipWriter.close()

    def zipDirectory(self, targetPath, localZipFilePath, createLocalPath=True):
        '''
        @see: transport.zipDirectory
        '''
        if not self.isDirectory(targetPath):
            error = 'Filepath %s is not a valid directory' % targetPath
            logger.warning(error)
            raise FileException(error)

        ensureFileCanBeCreated(localZipFilePath, createLocalPath,
                                     isFileExpected=True, isLocal=True)

        zipWriter = None
        try:
            zipWriter = ZipFile(localZipFilePath, 'w')

            logger.debug("Handle zipping a directory %s to archive %s",
                         targetPath, localZipFilePath)
            localZipFileFullPath = os.path.abspath(localZipFilePath)
            for root, _dirs, files in os.walk(targetPath):
                for f in files:
                    path = os.path.join(root, f)
                    # If current zip archive is located in this directory,
                    # exclude it from the archive
                    if os.path.abspath(path) != localZipFileFullPath:
                        zipWriter.write(path)
        finally:
            if zipWriter:
                zipWriter.close()

    def absPath(self, targetPath):
        '''
        @see: transport.absPath
        '''
        return os.path.abspath(targetPath)

    def startProcess(self, commandArgs, cwd=None, localStdinFile=None,
                    localStdoutFile=None, localStderrFile=None, env=None):
        '''
        @see: transport.startProcess
        '''
        if not isinstance(commandArgs, (list,tuple)):
            logger.warning("BAD REQUEST: Invalid exec command %s", commandArgs)
            raise ExecutionException("Invalid exec command %s" % commandArgs,
                                     ErrorCode.INVALID_ARGUMENTS)

        if len(commandArgs) == 0:
            logger.warning("BAD REQUEST: Invalid command %s", commandArgs)
            raise ExecutionException("Invalid command %s" % commandArgs)

        # Pre-process command and prepare for execution
        command = [os.path.expandvars(arg) for arg in commandArgs]

        stdin = None
        stdout = None
        stderr = None
        try:
            if localStdinFile:
                ensureFileExists(localStdinFile, isFileExpected=True, isLocal=True)
                stdin = open(localStdinFile)

            if localStdoutFile:
                ensureFileCanBeCreated(localStdoutFile, createPath=True,
                                             isFileExpected=True, isLocal=True)
                stdout = open(localStdoutFile, 'wb')
            else:
                stdout = tempfile.NamedTemporaryFile(prefix='stdout', delete=False)

            if localStderrFile:
                ensureFileCanBeCreated(localStderrFile, createPath=True,
                                             isFileExpected=True, isLocal=True)
                stderr = open(localStderrFile, 'wb')
            else:
                stderr = tempfile.NamedTemporaryFile(prefix='stderr', delete=False)

            close_fds = False if self._isWindowsPlatform() else True
            if env:
                # Extend local environment with environment given by the user
                localEnv = os.environ.copy()
                localEnv.update(env)
                env = localEnv

            # This will allow us to send CTRL+BREAK signals in case of
            # shutdown, to child processes
            creationflags = transport.CREATE_NEW_PROCESS_GROUP

            startedAt = time.time()
            popenHandle = subprocess.Popen(
                    command,
                    stdin=stdin,
                    stdout=stdout,
                    stderr=stderr,
                    cwd=cwd,
                    env=env,
                    close_fds=close_fds,
                    shell=False,
                    creationflags=creationflags)
        except OSError as e:
            # command does not exist
            logger.warning("BAD REQUEST: Cannot execute %s. Error: %s",
                        command, e)
            if e.errno == errno.ENOENT:
                error_code = ErrorCode.INVALID_REQUEST
            else:
                error_code = ErrorCode.UNKNOWN
            raise ExecutionException("Cannot execute %s. Error: %s" \
                                     % (command, e), error_code)
        finally:
            if stdin:
                stdin.close()
            if stdout:
                stdout.close()
            if stderr:
                stderr.close()

        proc = _Process(popenHandle, startedAt,
                        stdoutTempFile=None if localStdoutFile else stdout.name,
                        stderrTempFile=None if localStderrFile else stderr.name)

        self.allProcesses[proc.processUuid] = proc
        return proc.processUuid

    def pollProcess(self, processUuid, waitForCompletion=False, keepAsBytes=PY2):
        '''
        @see: transport.pollProcess
        '''
        if processUuid not in self.allProcesses:
            raise ExecutionException('Unknown process %s' % processUuid)

        proc = self.allProcesses[processUuid]

        result = proc.getProcessResult(waitForCompletion, keepAsBytes=keepAsBytes)

        # Remove the completed process and return its result
        if result.isCompleted():
            del self.allProcesses[processUuid]

        return result

    def getEnvironmentVariable(self, name, defaultValue=None):
        '''
        @see: transport.getEnvironmentVariable
        '''
        varValue = os.getenv(name, defaultValue)
        logger.debug("Environment variable '%s'=%s", name, varValue)
        return varValue

    def getPlatform(self):
        '''
        @see: transport.getPlatform
        '''
        return platform.system()

    def getAddress(self):
        '''
        @see: transport.getAdsress
        '''
        return 'localhost'

def _onShutdownRequest(allProcesses):
    '''Called on shutdown. The function goes and kill all runnning processes.

    @param allProcesses: All started processes.
    @type allProcess: Dictionary of Proceess Uuid to _LocalProcessResult
    '''
    for _uuid, processResult in allProcesses.iteritems():
        if not processResult.isCompleted():
            killProcess(processResult.getPid())

# Define shutdown hook. In case of shutdown signal go and kill all running processes
shutdownHookExecutor.addShutdownHook(_onShutdownRequest,
                                     LocalOperationsManager.allProcesses,
                                     shutdown_executor.ShutdownhookPriority.SYSTEM_HOOK)