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

"""This module controls locking and unlocking of test machines."""


import argparse
import enum
import getpass
import os
import sys

from cros_utils import command_executer
from cros_utils import logger
from cros_utils import machines
import file_lock_machine


class LockException(Exception):
    """Base class for exceptions in this module."""


class MachineNotPingable(LockException):
    """Raised when machine does not respond to ping."""


class LockingError(LockException):
    """Raised when server fails to lock/unlock machine as requested."""


class DontOwnLock(LockException):
    """Raised when user attmepts to unlock machine locked by someone else."""

    # This should not be raised if the user specified '--force'


class MachineType(enum.Enum):
    """Enum class to hold machine type."""

    LOCAL = "local"
    CROSFLEET = "crosfleet"


class LockManager(object):
    """Class for locking/unlocking machines vie three different modes.

    This class contains methods for checking the locked status of machines,
    and for changing the locked status.  It handles HW lab machines and local
    machines, using appropriate locking mechanisms for each.
    """

    CROSFLEET_PATH = "crosfleet"

    # TODO(zhizhouy): lease time may needs to be dynamically adjusted. For now we
    # set it long enough to cover the period to finish nightly rotation tests.
    LEASE_MINS = 1439

    CROSFLEET_CREDENTIAL = (
        "/usr/local/google/home/mobiletc-prebuild"
        "/sheriff_utils/credentials/skylab"
        "/chromeos-swarming-credential.json"
    )
    SWARMING = "~/cipd_binaries/swarming"
    SUCCESS = 0

    def __init__(
        self, remotes, force_option, chromeos_root, locks_dir="", log=None
    ):
        """Initializes an LockManager object.

        Args:
          remotes: A list of machine names or ip addresses to be managed.  Names
            and ip addresses should be represented as strings.  If the list is
            empty, the lock manager will get all known machines.
          force_option: A Boolean indicating whether or not to force an unlock of
            a machine that was locked by someone else.
          chromeos_root: The ChromeOS chroot to use for the autotest scripts.
          locks_dir: A directory used for file locking local devices.
          log: If not None, this is the logger object to be used for writing out
            informational output messages.  It is expected to be an instance of
            Logger class from cros_utils/logger.py.
        """
        self.chromeos_root = chromeos_root
        self.user = getpass.getuser()
        self.logger = log or logger.GetLogger()
        self.ce = command_executer.GetCommandExecuter(self.logger)

        sys.path.append(chromeos_root)

        self.locks_dir = locks_dir

        self.machines = list(set(remotes)) or []
        self.toolchain_lab_machines = self.GetAllToolchainLabMachines()

        if not self.machines:
            self.machines = self.toolchain_lab_machines
        self.force = force_option

        self.local_machines = []
        self.crosfleet_machines = []

    def CheckMachine(self, machine, error_msg):
        """Verifies that machine is responding to ping.

        Args:
          machine: String containing the name or ip address of machine to check.
          error_msg: Message to print if ping fails.

        Raises:
          MachineNotPingable:  If machine is not responding to 'ping'
        """
        if not machines.MachineIsPingable(machine, logging_level="none"):
            cros_machine = machine + ".cros"
            if not machines.MachineIsPingable(
                cros_machine, logging_level="none"
            ):
                raise MachineNotPingable(error_msg)

    def GetAllToolchainLabMachines(self):
        """Gets a list of all the toolchain machines in the ChromeOS HW lab.

        Returns:
          A list of names of the toolchain machines in the ChromeOS HW lab.
        """
        machines_file = os.path.join(
            os.path.dirname(__file__), "crosperf", "default_remotes"
        )
        machine_list = []
        with open(machines_file, "r") as input_file:
            lines = input_file.readlines()
            for line in lines:
                _, remotes = line.split(":")
                remotes = remotes.strip()
                for r in remotes.split():
                    machine_list.append(r.strip())
        return machine_list

    def GetMachineType(self, m):
        """Get where the machine is located.

        Args:
          m: String containing the name or ip address of machine.

        Returns:
          Value of the type in MachineType Enum.
        """
        if m in self.local_machines:
            return MachineType.LOCAL
        if m in self.crosfleet_machines:
            return MachineType.CROSFLEET

    def PrintStatusHeader(self):
        """Prints the status header lines for machines."""
        print("\nMachine (Board)\t\t\t\t\tStatus")
        print("---------------\t\t\t\t\t------")

    def PrintStatus(self, m, state, machine_type):
        """Prints status for a single machine.

        Args:
          m: String containing the name or ip address of machine.
          state: A dictionary of the current state of the machine.
          machine_type: MachineType to determine where the machine is located.
        """
        if state["locked"]:
            print(
                "%s (%s)\t\t%slocked by %s since %s"
                % (
                    m,
                    state["board"],
                    "\t\t" if machine_type == MachineType.LOCAL else "",
                    state["locked_by"],
                    state["lock_time"],
                )
            )
        else:
            print(
                "%s (%s)\t\t%sunlocked"
                % (
                    m,
                    state["board"],
                    "\t\t" if machine_type == MachineType.LOCAL else "",
                )
            )

    def AddMachineToLocal(self, machine):
        """Adds a machine to local machine list.

        Args:
          machine: The machine to be added.
        """
        if machine not in self.local_machines:
            self.local_machines.append(machine)

    def AddMachineToCrosfleet(self, machine):
        """Adds a machine to crosfleet machine list.

        Args:
          machine: The machine to be added.
        """
        if machine not in self.crosfleet_machines:
            self.crosfleet_machines.append(machine)

    def ListMachineStates(self, machine_states):
        """Gets and prints the current status for a list of machines.

        Prints out the current status for all of the machines in the current
        LockManager's list of machines (set when the object is initialized).

        Args:
          machine_states: A dictionary of the current state of every machine in
            the current LockManager's list of machines.  Normally obtained by
            calling LockManager::GetMachineStates.
        """
        self.PrintStatusHeader()
        for m in machine_states:
            machine_type = self.GetMachineType(m)
            state = machine_states[m]
            self.PrintStatus(m, state, machine_type)

    def UpdateLockInCrosfleet(self, should_lock_machine, machine):
        """Ask crosfleet to lease/release a machine.

        Args:
          should_lock_machine: Boolean indicating whether to lock the machine (True)
            or unlock the machine (False).
          machine: The machine to update.

        Returns:
          True if requested action succeeded, else False.
        """
        try:
            if should_lock_machine:
                ret = self.LeaseCrosfleetMachine(machine)
            else:
                ret = self.ReleaseCrosfleetMachine(machine)
        except Exception:
            return False
        return ret

    def UpdateFileLock(self, should_lock_machine, machine):
        """Use file lock for local machines,

        Args:
          should_lock_machine: Boolean indicating whether to lock the machine (True)
            or unlock the machine (False).
          machine: The machine to update.

        Returns:
          True if requested action succeeded, else False.
        """
        try:
            if should_lock_machine:
                ret = file_lock_machine.Machine(machine, self.locks_dir).Lock(
                    True, sys.argv[0]
                )
            else:
                ret = file_lock_machine.Machine(machine, self.locks_dir).Unlock(
                    True
                )
        except Exception:
            return False
        return ret

    def UpdateMachines(self, lock_machines):
        """Sets the locked state of the machines to the requested value.

        The machines updated are the ones in self.machines (specified when the
        class object was intialized).

        Args:
          lock_machines: Boolean indicating whether to lock the machines (True) or
            unlock the machines (False).

        Returns:
          A list of the machines whose state was successfully updated.
        """
        updated_machines = []
        action = "Locking" if lock_machines else "Unlocking"
        for m in self.machines:
            # TODO(zhizhouy): Handling exceptions with more details when locking
            # doesn't succeed.
            machine_type = self.GetMachineType(m)
            if machine_type == MachineType.CROSFLEET:
                ret = self.UpdateLockInCrosfleet(lock_machines, m)
            elif machine_type == MachineType.LOCAL:
                ret = self.UpdateFileLock(lock_machines, m)

            if ret:
                self.logger.LogOutput(
                    "%s %s machine succeeded: %s."
                    % (action, machine_type.value, m)
                )
                updated_machines.append(m)
            else:
                self.logger.LogOutput(
                    "%s %s machine failed: %s."
                    % (action, machine_type.value, m)
                )

        self.machines = updated_machines
        return updated_machines

    def _InternalRemoveMachine(self, machine):
        """Remove machine from internal list of machines.

        Args:
          machine: Name of machine to be removed from internal list.
        """
        # Check to see if machine is lab machine and if so, make sure it has
        # ".cros" on the end.
        cros_machine = machine
        if machine.find("rack") > 0 and machine.find("row") > 0:
            if machine.find(".cros") == -1:
                cros_machine = cros_machine + ".cros"

        self.machines = [
            m for m in self.machines if m not in (cros_machine, machine)
        ]

    def CheckMachineLocks(self, machine_states, cmd):
        """Check that every machine in requested list is in the proper state.

        If the cmd is 'unlock' verify that every machine is locked by requestor.
        If the cmd is 'lock' verify that every machine is currently unlocked.

        Args:
          machine_states: A dictionary of the current state of every machine in
            the current LockManager's list of machines.  Normally obtained by
            calling LockManager::GetMachineStates.
          cmd: The user-requested action for the machines: 'lock' or 'unlock'.

        Raises:
          DontOwnLock: The lock on a requested machine is owned by someone else.
        """
        for k, state in machine_states.items():
            if cmd == "unlock":
                if not state["locked"]:
                    self.logger.LogWarning(
                        "Attempt to unlock already unlocked machine "
                        "(%s)." % k
                    )
                    self._InternalRemoveMachine(k)

                # TODO(zhizhouy): Crosfleet doesn't support host info such as locked_by.
                # Need to update this when crosfleet supports it.
                if (
                    state["locked"]
                    and state["locked_by"]
                    and state["locked_by"] != self.user
                ):
                    raise DontOwnLock(
                        "Attempt to unlock machine (%s) locked by someone "
                        "else (%s)." % (k, state["locked_by"])
                    )
            elif cmd == "lock":
                if state["locked"]:
                    self.logger.LogWarning(
                        "Attempt to lock already locked machine (%s)" % k
                    )
                    self._InternalRemoveMachine(k)

    def GetMachineStates(self, cmd=""):
        """Gets the current state of all the requested machines.

        Gets the current state of all the requested machines. Stores the data in a
        dictionary keyed by machine name.

        Args:
          cmd: The command for which we are getting the machine states. This is
            important because if one of the requested machines is missing we raise
            an exception, unless the requested command is 'add'.

        Returns:
          A dictionary of machine states for all the machines in the LockManager
          object.
        """
        machine_list = {}
        for m in self.machines:
            # For local or crosfleet machines, we simply set {'locked': status} for
            # them
            # TODO(zhizhouy): This is a quick fix since crosfleet cannot return host
            # info as afe does. We need to get more info such as locked_by when
            # crosfleet supports that.
            values = {
                "locked": 0 if cmd == "lock" else 1,
                "board": "??",
                "locked_by": "",
                "lock_time": "",
            }
            machine_list[m] = values

        self.ListMachineStates(machine_list)

        return machine_list

    def CheckMachineInCrosfleet(self, machine):
        """Run command to check if machine is in Crosfleet or not.

        Returns:
          True if machine in crosfleet, else False
        """
        credential = ""
        if os.path.exists(self.CROSFLEET_CREDENTIAL):
            credential = "--service-account-json %s" % self.CROSFLEET_CREDENTIAL
        server = "--server https://chromeos-swarming.appspot.com"
        dimensions = "--dimension dut_name=%s" % machine.rstrip(".cros")

        cmd = f"{self.SWARMING} bots {server} {credential} {dimensions}"
        exit_code, stdout, stderr = self.ce.RunCommandWOutput(cmd)
        if exit_code:
            raise ValueError(
                "Querying bots failed (2); stdout: %r; stderr: %r"
                % (stdout, stderr)
            )

        # The command will return a json output as stdout. If machine not in
        # crosfleet, stdout will look like this:
        #  {
        #    "death_timeout": "600",
        #    "now": "TIMESTAMP"
        #  }
        # Otherwise there will be a tuple starting with 'items', we simply detect
        # this keyword for result.
        return stdout != "[]"

    def LeaseCrosfleetMachine(self, machine):
        """Run command to lease dut from crosfleet.

        Returns:
          True if succeeded, False if failed.
        """
        credential = ""
        if os.path.exists(self.CROSFLEET_CREDENTIAL):
            credential = "-service-account-json %s" % self.CROSFLEET_CREDENTIAL
        cmd = ("%s dut lease -minutes %s %s %s %s") % (
            self.CROSFLEET_PATH,
            self.LEASE_MINS,
            credential,
            "-host",
            machine.rstrip(".cros"),
        )
        # Wait 8 minutes for server to start the lease task, if not started,
        # we will treat it as unavailable.
        check_interval_time = 480
        retval = self.ce.RunCommand(cmd, command_timeout=check_interval_time)
        return retval == self.SUCCESS

    def ReleaseCrosfleetMachine(self, machine):
        """Run command to release dut from crosfleet.

        Returns:
          True if succeeded, False if failed.
        """
        credential = ""
        if os.path.exists(self.CROSFLEET_CREDENTIAL):
            credential = "-service-account-json %s" % self.CROSFLEET_CREDENTIAL

        cmd = ("%s dut abandon %s %s") % (
            self.CROSFLEET_PATH,
            credential,
            machine.rstrip(".cros"),
        )
        retval = self.ce.RunCommand(cmd)
        return retval == self.SUCCESS


def Main(argv):
    """Parse the options, initialize lock manager and dispatch proper method.

    Args:
      argv: The options with which this script was invoked.

    Returns:
      0 unless an exception is raised.
    """
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--list",
        dest="cmd",
        action="store_const",
        const="status",
        help="List current status of all known machines.",
    )
    parser.add_argument(
        "--lock",
        dest="cmd",
        action="store_const",
        const="lock",
        help="Lock given machine(s).",
    )
    parser.add_argument(
        "--unlock",
        dest="cmd",
        action="store_const",
        const="unlock",
        help="Unlock given machine(s).",
    )
    parser.add_argument(
        "--status",
        dest="cmd",
        action="store_const",
        const="status",
        help="List current status of given machine(s).",
    )
    parser.add_argument(
        "--remote", dest="remote", help="machines on which to operate"
    )
    parser.add_argument(
        "--chromeos_root",
        dest="chromeos_root",
        required=True,
        help="ChromeOS root to use for autotest scripts.",
    )
    parser.add_argument(
        "--force",
        dest="force",
        action="store_true",
        default=False,
        help="Force lock/unlock of machines, even if not"
        " current lock owner.",
    )

    options = parser.parse_args(argv)

    if not options.remote and options.cmd != "status":
        parser.error("No machines specified for operation.")

    if not os.path.isdir(options.chromeos_root):
        parser.error("Cannot find chromeos_root: %s." % options.chromeos_root)

    if not options.cmd:
        parser.error(
            "No operation selected (--list, --status, --lock, --unlock,"
            " --add_machine, --remove_machine)."
        )

    machine_list = []
    if options.remote:
        machine_list = options.remote.split()

    lock_manager = LockManager(
        machine_list, options.force, options.chromeos_root
    )

    machine_states = lock_manager.GetMachineStates(cmd=options.cmd)
    cmd = options.cmd

    if cmd == "status":
        lock_manager.ListMachineStates(machine_states)

    elif cmd == "lock":
        if not lock_manager.force:
            lock_manager.CheckMachineLocks(machine_states, cmd)
            lock_manager.UpdateMachines(True)

    elif cmd == "unlock":
        if not lock_manager.force:
            lock_manager.CheckMachineLocks(machine_states, cmd)
            lock_manager.UpdateMachines(False)

    return 0


if __name__ == "__main__":
    sys.exit(Main(sys.argv[1:]))