#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2011 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Utilities to run commands in outside/inside chroot and on the board."""


import getpass
import os
import re
import select
import signal
import subprocess
import sys
import tempfile
import time

from cros_utils import logger


mock_default = False

CHROMEOS_SCRIPTS_DIR = "/mnt/host/source/src/scripts"
LOG_LEVEL = ("none", "quiet", "average", "verbose")


def InitCommandExecuter(mock=False):
    # pylint: disable=global-statement
    global mock_default
    # Whether to default to a mock command executer or not
    mock_default = mock


def GetCommandExecuter(logger_to_set=None, mock=False, log_level="verbose"):
    # If the default is a mock executer, always return one.
    if mock_default or mock:
        return MockCommandExecuter(log_level, logger_to_set)
    else:
        return CommandExecuter(log_level, logger_to_set)


class CommandExecuter(object):
    """Provides several methods to execute commands on several environments."""

    def __init__(self, log_level, logger_to_set=None):
        self.log_level = log_level
        if log_level == "none":
            self.logger = None
        else:
            if logger_to_set is not None:
                self.logger = logger_to_set
            else:
                self.logger = logger.GetLogger()

    def GetLogLevel(self):
        return self.log_level

    def SetLogLevel(self, log_level):
        self.log_level = log_level

    def RunCommandGeneric(
        self,
        cmd,
        return_output=False,
        machine=None,
        username=None,
        command_terminator=None,
        command_timeout=None,
        terminated_timeout=10,
        print_to_console=True,
        env=None,
        except_handler=lambda p, e: None,
    ):
        """Run a command.

        Returns triplet (returncode, stdout, stderr).
        """

        cmd = str(cmd)

        if self.log_level == "quiet":
            print_to_console = False

        if self.log_level == "verbose":
            self.logger.LogCmd(cmd, machine, username, print_to_console)
        elif self.logger:
            self.logger.LogCmdToFileOnly(cmd, machine, username)
        if command_terminator and command_terminator.IsTerminated():
            if self.logger:
                self.logger.LogError(
                    "Command was terminated!", print_to_console
                )
            return (1, "", "")

        if machine is not None:
            user = ""
            if username is not None:
                user = username + "@"
            cmd = "ssh -t -t %s%s -- '%s'" % (user, machine, cmd)

        # We use setsid so that the child will have a different session id
        # and we can easily kill the process group. This is also important
        # because the child will be disassociated from the parent terminal.
        # In this way the child cannot mess the parent's terminal.
        p = None
        try:
            # pylint: disable=bad-option-value, subprocess-popen-preexec-fn
            p = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                shell=True,
                preexec_fn=os.setsid,
                executable="/bin/bash",
                env=env,
            )

            full_stdout = ""
            full_stderr = ""

            # Pull output from pipes, send it to file/stdout/string
            out = err = None
            pipes = [p.stdout, p.stderr]

            my_poll = select.poll()
            my_poll.register(p.stdout, select.POLLIN)
            my_poll.register(p.stderr, select.POLLIN)

            terminated_time = None
            started_time = time.time()

            while pipes:
                if command_terminator and command_terminator.IsTerminated():
                    os.killpg(os.getpgid(p.pid), signal.SIGTERM)
                    if self.logger:
                        self.logger.LogError(
                            "Command received termination request. "
                            "Killed child process group.",
                            print_to_console,
                        )
                    break

                l = my_poll.poll(100)
                for (fd, _) in l:
                    if fd == p.stdout.fileno():
                        out = os.read(p.stdout.fileno(), 16384).decode("utf8")
                        if return_output:
                            full_stdout += out
                        if self.logger:
                            self.logger.LogCommandOutput(out, print_to_console)
                        if out == "":
                            pipes.remove(p.stdout)
                            my_poll.unregister(p.stdout)
                    if fd == p.stderr.fileno():
                        err = os.read(p.stderr.fileno(), 16384).decode("utf8")
                        if return_output:
                            full_stderr += err
                        if self.logger:
                            self.logger.LogCommandError(err, print_to_console)
                        if err == "":
                            pipes.remove(p.stderr)
                            my_poll.unregister(p.stderr)

                if p.poll() is not None:
                    if terminated_time is None:
                        terminated_time = time.time()
                    elif (
                        terminated_timeout is not None
                        and time.time() - terminated_time > terminated_timeout
                    ):
                        if self.logger:
                            self.logger.LogWarning(
                                "Timeout of %s seconds reached since "
                                "process termination." % terminated_timeout,
                                print_to_console,
                            )
                        break

                if (
                    command_timeout is not None
                    and time.time() - started_time > command_timeout
                ):
                    os.killpg(os.getpgid(p.pid), signal.SIGTERM)
                    if self.logger:
                        self.logger.LogWarning(
                            "Timeout of %s seconds reached since process"
                            "started. Killed child process group."
                            % command_timeout,
                            print_to_console,
                        )
                    break

                if out == err == "":
                    break

            p.wait()
            if return_output:
                return (p.returncode, full_stdout, full_stderr)
            return (p.returncode, "", "")
        except BaseException as err:
            except_handler(p, err)
            raise

    def RunCommand(self, *args, **kwargs):
        """Run a command.

        Takes the same arguments as RunCommandGeneric except for return_output.
        Returns a single value returncode.
        """
        # Make sure that args does not overwrite 'return_output'
        assert len(args) <= 1
        assert "return_output" not in kwargs
        kwargs["return_output"] = False
        return self.RunCommandGeneric(*args, **kwargs)[0]

    def RunCommandWExceptionCleanup(self, *args, **kwargs):
        """Run a command and kill process if exception is thrown.

        Takes the same arguments as RunCommandGeneric except for except_handler.
        Returns same as RunCommandGeneric.
        """

        def KillProc(proc, _):
            if proc:
                os.killpg(os.getpgid(proc.pid), signal.SIGTERM)

        # Make sure that args does not overwrite 'except_handler'
        assert len(args) <= 8
        assert "except_handler" not in kwargs
        kwargs["except_handler"] = KillProc
        return self.RunCommandGeneric(*args, **kwargs)

    def RunCommandWOutput(self, *args, **kwargs):
        """Run a command.

        Takes the same arguments as RunCommandGeneric except for return_output.
        Returns a triplet (returncode, stdout, stderr).
        """
        # Make sure that args does not overwrite 'return_output'
        assert len(args) <= 1
        assert "return_output" not in kwargs
        kwargs["return_output"] = True
        return self.RunCommandGeneric(*args, **kwargs)

    def RemoteAccessInitCommand(self, chromeos_root, machine, port=None):
        command = ""
        command += "\nset -- --remote=" + machine
        if port:
            command += " --ssh_port=" + port
        command += "\n. " + chromeos_root + "/src/scripts/common.sh"
        command += "\n. " + chromeos_root + "/src/scripts/remote_access.sh"
        command += "\nTMP=$(mktemp -d)"
        command += '\nFLAGS "$@" || exit 1'
        command += "\nremote_access_init"
        return command

    def WriteToTempShFile(self, contents):
        with tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            delete=False,
            prefix=os.uname()[1],
            suffix=".sh",
        ) as f:
            f.write("#!/bin/bash\n")
            f.write(contents)
            f.flush()
        return f.name

    def CrosLearnBoard(self, chromeos_root, machine):
        command = self.RemoteAccessInitCommand(chromeos_root, machine)
        command += "\nlearn_board"
        command += "\necho ${FLAGS_board}"
        retval, output, _ = self.RunCommandWOutput(command)
        if self.logger:
            self.logger.LogFatalIf(retval, "learn_board command failed")
        elif retval:
            sys.exit(1)
        return output.split()[-1]

    def CrosRunCommandGeneric(
        self,
        cmd,
        return_output=False,
        machine=None,
        command_terminator=None,
        chromeos_root=None,
        command_timeout=None,
        terminated_timeout=10,
        print_to_console=True,
    ):
        """Run a command on a ChromeOS box.

        Returns triplet (returncode, stdout, stderr).
        """

        if self.log_level != "verbose":
            print_to_console = False

        if self.logger:
            self.logger.LogCmd(cmd, print_to_console=print_to_console)
            self.logger.LogFatalIf(not machine, "No machine provided!")
            self.logger.LogFatalIf(
                not chromeos_root, "chromeos_root not given!"
            )
        else:
            if not chromeos_root or not machine:
                sys.exit(1)
        chromeos_root = os.path.expanduser(chromeos_root)

        port = None
        if ":" in machine:
            machine, port = machine.split(":")
        # Write all commands to a file.
        command_file = self.WriteToTempShFile(cmd)
        retval = self.CopyFiles(
            command_file,
            command_file,
            dest_machine=machine,
            dest_port=port,
            command_terminator=command_terminator,
            chromeos_root=chromeos_root,
            dest_cros=True,
            recursive=False,
            print_to_console=print_to_console,
        )
        if retval:
            if self.logger:
                self.logger.LogError(
                    "Could not run remote command on machine."
                    " Is the machine up?"
                )
            return (retval, "", "")

        command = self.RemoteAccessInitCommand(chromeos_root, machine, port)
        command += "\nremote_sh bash %s" % command_file
        command += '\nl_retval=$?; echo "$REMOTE_OUT"; exit $l_retval'
        retval = self.RunCommandGeneric(
            command,
            return_output,
            command_terminator=command_terminator,
            command_timeout=command_timeout,
            terminated_timeout=terminated_timeout,
            print_to_console=print_to_console,
        )
        if return_output:
            connect_signature = (
                "Initiating first contact with remote host\n"
                + "Connection OK\n"
            )
            connect_signature_re = re.compile(connect_signature)
            modded_retval = list(retval)
            modded_retval[1] = connect_signature_re.sub("", retval[1])
            return modded_retval
        return retval

    def CrosRunCommand(self, *args, **kwargs):
        """Run a command on a ChromeOS box.

        Takes the same arguments as CrosRunCommandGeneric except for return_output.
        Returns a single value returncode.
        """
        # Make sure that args does not overwrite 'return_output'
        assert len(args) <= 1
        assert "return_output" not in kwargs
        kwargs["return_output"] = False
        return self.CrosRunCommandGeneric(*args, **kwargs)[0]

    def CrosRunCommandWOutput(self, *args, **kwargs):
        """Run a command on a ChromeOS box.

        Takes the same arguments as CrosRunCommandGeneric except for return_output.
        Returns a triplet (returncode, stdout, stderr).
        """
        # Make sure that args does not overwrite 'return_output'
        assert len(args) <= 1
        assert "return_output" not in kwargs
        kwargs["return_output"] = True
        return self.CrosRunCommandGeneric(*args, **kwargs)

    def ChrootRunCommandGeneric(
        self,
        chromeos_root,
        command,
        return_output=False,
        command_terminator=None,
        command_timeout=None,
        terminated_timeout=10,
        print_to_console=True,
        cros_sdk_options="",
        env=None,
    ):
        """Runs a command within the chroot.

        Returns triplet (returncode, stdout, stderr).
        """

        if self.log_level != "verbose":
            print_to_console = False

        if self.logger:
            self.logger.LogCmd(command, print_to_console=print_to_console)

        with tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            delete=False,
            dir=os.path.join(chromeos_root, "src/scripts"),
            suffix=".sh",
            prefix="in_chroot_cmd",
        ) as f:
            f.write("#!/bin/bash\n")
            f.write(command)
            f.write("\n")
            f.flush()

        command_file = f.name
        os.chmod(command_file, 0o777)

        # if return_output is set, run a test command first to make sure that
        # the chroot already exists. We want the final returned output to skip
        # the output from chroot creation steps.
        if return_output:
            ret = self.RunCommand(
                "cd %s; cros_sdk %s -- true"
                % (chromeos_root, cros_sdk_options),
                env=env,
                # Give this command a long time to execute; it might involve setting
                # the chroot up, or running fstrim on its image file. Both of these
                # operations can take well over the timeout default of 10 seconds.
                terminated_timeout=5 * 60,
            )
            if ret:
                return (ret, "", "")

        # Run command_file inside the chroot, making sure that any "~" is expanded
        # by the shell inside the chroot, not outside.
        command = "cd %s; cros_sdk %s -- bash -c '%s/%s'" % (
            chromeos_root,
            cros_sdk_options,
            CHROMEOS_SCRIPTS_DIR,
            os.path.basename(command_file),
        )
        ret = self.RunCommandGeneric(
            command,
            return_output,
            command_terminator=command_terminator,
            command_timeout=command_timeout,
            terminated_timeout=terminated_timeout,
            print_to_console=print_to_console,
            env=env,
        )
        os.remove(command_file)
        return ret

    def ChrootRunCommand(self, *args, **kwargs):
        """Runs a command within the chroot.

        Takes the same arguments as ChrootRunCommandGeneric except for
        return_output.
        Returns a single value returncode.
        """
        # Make sure that args does not overwrite 'return_output'
        assert len(args) <= 2
        assert "return_output" not in kwargs
        kwargs["return_output"] = False
        return self.ChrootRunCommandGeneric(*args, **kwargs)[0]

    def ChrootRunCommandWOutput(self, *args, **kwargs):
        """Runs a command within the chroot.

        Takes the same arguments as ChrootRunCommandGeneric except for
        return_output.
        Returns a triplet (returncode, stdout, stderr).
        """
        # Make sure that args does not overwrite 'return_output'
        assert len(args) <= 2
        assert "return_output" not in kwargs
        kwargs["return_output"] = True
        return self.ChrootRunCommandGeneric(*args, **kwargs)

    def RunCommands(
        self, cmdlist, machine=None, username=None, command_terminator=None
    ):
        cmd = " ;\n".join(cmdlist)
        return self.RunCommand(
            cmd,
            machine=machine,
            username=username,
            command_terminator=command_terminator,
        )

    def CopyFiles(
        self,
        src,
        dest,
        src_machine=None,
        src_port=None,
        dest_machine=None,
        dest_port=None,
        src_user=None,
        dest_user=None,
        recursive=True,
        command_terminator=None,
        chromeos_root=None,
        src_cros=False,
        dest_cros=False,
        print_to_console=True,
    ):
        src = os.path.expanduser(src)
        dest = os.path.expanduser(dest)

        if recursive:
            src = src + "/"
            dest = dest + "/"

        if src_cros or dest_cros:
            if self.logger:
                self.logger.LogFatalIf(
                    src_cros == dest_cros,
                    "Only one of src_cros and desc_cros can " "be True.",
                )
                self.logger.LogFatalIf(
                    not chromeos_root, "chromeos_root not given!"
                )
            elif src_cros == dest_cros or not chromeos_root:
                sys.exit(1)
            if src_cros:
                cros_machine = src_machine
                cros_port = src_port
                host_machine = dest_machine
                host_user = dest_user
            else:
                cros_machine = dest_machine
                cros_port = dest_port
                host_machine = src_machine
                host_user = src_user

            command = self.RemoteAccessInitCommand(
                chromeos_root, cros_machine, cros_port
            )
            ssh_command = (
                "ssh -o StrictHostKeyChecking=no"
                + " -o UserKnownHostsFile=$(mktemp)"
                + " -i $TMP_PRIVATE_KEY"
            )
            if cros_port:
                ssh_command += " -p %s" % cros_port
            rsync_prefix = '\nrsync -r -e "%s" ' % ssh_command
            if dest_cros:
                command += rsync_prefix + "%s root@%s:%s" % (
                    src,
                    cros_machine,
                    dest,
                )
            else:
                command += rsync_prefix + "root@%s:%s %s" % (
                    cros_machine,
                    src,
                    dest,
                )

            return self.RunCommand(
                command,
                machine=host_machine,
                username=host_user,
                command_terminator=command_terminator,
                print_to_console=print_to_console,
            )

        if dest_machine == src_machine:
            command = "rsync -a %s %s" % (src, dest)
        else:
            if src_machine is None:
                src_machine = os.uname()[1]
                src_user = getpass.getuser()
            command = "rsync -a %s@%s:%s %s" % (
                src_user,
                src_machine,
                src,
                dest,
            )
        return self.RunCommand(
            command,
            machine=dest_machine,
            username=dest_user,
            command_terminator=command_terminator,
            print_to_console=print_to_console,
        )

    def RunCommand2(
        self,
        cmd,
        cwd=None,
        line_consumer=None,
        timeout=None,
        shell=True,
        join_stderr=True,
        env=None,
        except_handler=lambda p, e: None,
    ):
        """Run the command with an extra feature line_consumer.

        This version allow developers to provide a line_consumer which will be
        fed execution output lines.

        A line_consumer is a callback, which is given a chance to run for each
        line the execution outputs (either to stdout or stderr). The
        line_consumer must accept one and exactly one dict argument, the dict
        argument has these items -
          'line'   -  The line output by the binary. Notice, this string includes
                      the trailing '\n'.
          'output' -  Whether this is a stdout or stderr output, values are either
                      'stdout' or 'stderr'. When join_stderr is True, this value
                      will always be 'output'.
          'pobject' - The object used to control execution, for example, call
                      pobject.kill().

        Note: As this is written, the stdin for the process executed is
        not associated with the stdin of the caller of this routine.

        Args:
          cmd: Command in a single string.
          cwd: Working directory for execution.
          line_consumer: A function that will ba called by this function. See above
            for details.
          timeout: terminate command after this timeout.
          shell: Whether to use a shell for execution.
          join_stderr: Whether join stderr to stdout stream.
          env: Execution environment.
          except_handler: Callback for when exception is thrown during command
            execution. Passed process object and exception.

        Returns:
          Execution return code.

        Raises:
          child_exception: if fails to start the command process (missing
                           permission, no such file, etc)
        """

        class StreamHandler(object):
            """Internal utility class."""

            def __init__(self, pobject, fd, name, line_consumer):
                self._pobject = pobject
                self._fd = fd
                self._name = name
                self._buf = ""
                self._line_consumer = line_consumer

            def read_and_notify_line(self):
                t = os.read(fd, 1024)
                self._buf = self._buf + t
                self.notify_line()

            def notify_line(self):
                p = self._buf.find("\n")
                while p >= 0:
                    self._line_consumer(
                        line=self._buf[: p + 1],
                        output=self._name,
                        pobject=self._pobject,
                    )
                    if p < len(self._buf) - 1:
                        self._buf = self._buf[p + 1 :]
                        p = self._buf.find("\n")
                    else:
                        self._buf = ""
                        p = -1
                        break

            def notify_eos(self):
                # Notify end of stream. The last line may not end with a '\n'.
                if self._buf != "":
                    self._line_consumer(
                        line=self._buf, output=self._name, pobject=self._pobject
                    )
                    self._buf = ""

        if self.log_level == "verbose":
            self.logger.LogCmd(cmd)
        elif self.logger:
            self.logger.LogCmdToFileOnly(cmd)

        # We use setsid so that the child will have a different session id
        # and we can easily kill the process group. This is also important
        # because the child will be disassociated from the parent terminal.
        # In this way the child cannot mess the parent's terminal.
        pobject = None
        try:
            # pylint: disable=bad-option-value, subprocess-popen-preexec-fn
            pobject = subprocess.Popen(
                cmd,
                cwd=cwd,
                bufsize=1024,
                env=env,
                shell=shell,
                universal_newlines=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT if join_stderr else subprocess.PIPE,
                preexec_fn=os.setsid,
            )

            # We provide a default line_consumer
            if line_consumer is None:
                line_consumer = lambda **d: None
            start_time = time.time()
            poll = select.poll()
            outfd = pobject.stdout.fileno()
            poll.register(outfd, select.POLLIN | select.POLLPRI)
            handlermap = {
                outfd: StreamHandler(pobject, outfd, "stdout", line_consumer)
            }
            if not join_stderr:
                errfd = pobject.stderr.fileno()
                poll.register(errfd, select.POLLIN | select.POLLPRI)
                handlermap[errfd] = StreamHandler(
                    pobject, errfd, "stderr", line_consumer
                )
            while handlermap:
                readables = poll.poll(300)
                for (fd, evt) in readables:
                    handler = handlermap[fd]
                    if evt & (select.POLLPRI | select.POLLIN):
                        handler.read_and_notify_line()
                    elif evt & (
                        select.POLLHUP | select.POLLERR | select.POLLNVAL
                    ):
                        handler.notify_eos()
                        poll.unregister(fd)
                        del handlermap[fd]

                if timeout is not None and (time.time() - start_time > timeout):
                    os.killpg(os.getpgid(pobject.pid), signal.SIGTERM)

            return pobject.wait()
        except BaseException as err:
            except_handler(pobject, err)
            raise


class MockCommandExecuter(CommandExecuter):
    """Mock class for class CommandExecuter."""

    def RunCommandGeneric(
        self,
        cmd,
        return_output=False,
        machine=None,
        username=None,
        command_terminator=None,
        command_timeout=None,
        terminated_timeout=10,
        print_to_console=True,
        env=None,
        except_handler=lambda p, e: None,
    ):
        assert not command_timeout
        cmd = str(cmd)
        if machine is None:
            machine = "localhost"
        if username is None:
            username = "current"
        logger.GetLogger().LogCmd(
            "(Mock) " + cmd, machine, username, print_to_console
        )
        return (0, "", "")

    def RunCommand(self, *args, **kwargs):
        assert "return_output" not in kwargs
        kwargs["return_output"] = False
        return self.RunCommandGeneric(*args, **kwargs)[0]

    def RunCommandWOutput(self, *args, **kwargs):
        assert "return_output" not in kwargs
        kwargs["return_output"] = True
        return self.RunCommandGeneric(*args, **kwargs)


class CommandTerminator(object):
    """Object to request termination of a command in execution."""

    def __init__(self):
        self.terminated = False

    def Terminate(self):
        self.terminated = True

    def IsTerminated(self):
        return self.terminated