#!/usr/bin/env python
# Copyright 2018 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Implementation of graphyte.link.DeviceLink using SSH."""

import os
import pipes
import platform
import shlex
import subprocess
import threading
import time

import six
from six.moves import queue

import graphyte_common  # pylint: disable=unused-import
from graphyte import link
from graphyte.default_setting import logger
from graphyte.utils.file_utils import UnopenedTemporaryFile
from graphyte.utils import graphyte_utils
from graphyte.utils import process_utils
from graphyte.utils import sync_utils
from graphyte.utils import type_utils


class ClientNotExistError(Exception):
  def __str__(self):
    return 'There is no DHCP client registered.'


class SSHLink(link.DeviceLink):
  """A DUT target that is connected via SSH interface.

  Properties:
    host: A string for SSH host IP.
    user: A string for the user accont to login. Defaults to 'root'.
    port: An integer for the SSH port on remote host.
    identify: An identity file to specify credential.
    connect_timeout: An integer for ssh(1) connection timeout in seconds.
    control_persist: An integer for ssh(1) to keep master connection remain
              opened for given seconds, or None to not using master control.

  link_options example:
    {
      'link_class': 'SSHLink',
      'host': '1.2.3.4',
      'identity': '/path/to/identity/file'
    }
  """

  def __init__(self, host, user='root', port=22, identity=None,
               connect_timeout=5, control_persist=300):
    self.host = host
    self.user = user
    self.port = port
    self.identity = identity
    self.connect_timeout = connect_timeout
    self.control_persist = control_persist

    if self.identity is not None:
      try:
        self.identity = graphyte_utils.SearchConfig(self.identity)
        os.chmod(self.identity, 0o600)
      except Exception:
        logger.error('SSHLink: Failed to locate identity file "%s"',
                     self.identity)
        self.identity = None

  def __str__(self):
    return "SSH to %s@%s:%s with identity='%s'" % (
        self.user, self.host, self.port, self.identity)

  def _Signature(self, is_scp=False):
    """Generates the ssh command signature.

    Args:
      is_scp: A boolean flag indicating if the signature is made for scp.

    Returns:
      A pair of signature in (sig, options). The 'sig' is a string representing
      remote ssh user and host. 'options' is a list of required command line
      parameters.
    """
    if self.user:
      sig = '%s@%s' % (self.user, self.host)
    else:
      sig = self.host

    options = ['-o', 'UserKnownHostsFile=/dev/null',
               '-o', 'StrictHostKeyChecking=no',
               '-o', 'ConnectTimeout=%d' % self.connect_timeout]
    if self.control_persist is not None:
      options += ['-o', 'ControlMaster=auto',
                  '-o', 'ControlPath=/tmp/.ssh-%r@%h:%p',
                  '-o', 'ControlPersist=%s' % self.control_persist]
    if self.port:
      options += ['-P' if is_scp else '-p', str(self.port)]
    if self.identity:
      options += ['-i', self.identity]
    return sig, options

  def _ToUnixPath(self, path):
    if platform.system() != 'Windows':
      return path  # assume that we are in Unix machine already

    path = os.path.abspath(path)
    drive, path = os.path.splitdrive(path)
    return '/' + drive[0] + path.replace(os.path.sep, '/')

  def Push(self, local, remote):
    """See DUTLink.Push"""
    remote_sig, options = self._Signature(True)
    local = self._ToUnixPath(local)
    return subprocess.check_call(['scp'] + options +
                                 [local, '%s:%s' % (remote_sig, remote)])

  def Pull(self, remote, local=None):
    """See DUTLink.Pull"""
    if local is None:
      with UnopenedTemporaryFile() as path:
        self.Pull(remote, path)
        with open(path) as f:
          return f.read()

    local = self._ToUnixPath(local)
    remote_sig, options = self._Signature(True)
    subprocess.check_call(['scp'] + options +
                          ['%s:%s' % (remote_sig, remote), local])

  def _StartWatcher(self, subproc):
    watcher = self.__class__.ControlMasterWatcher(self)
    watcher.Start()  # make sure the watcher is running
    watcher.AddProcess(subproc.pid, os.getpid())

  def Shell(self, command, stdin=None, stdout=None, stderr=None):
    """See DUTLink.Shell"""
    remote_sig, options = self._Signature(False)

    if isinstance(command, six.string_types):
      # since the shell command is passed through SSH client,
      # the entire command is argument for SSH command
      # e.g. command = 'VAR=xxx cmd arg0 arg1...'
      # will become: 'ssh' <SSH OPTIONS> 'VAR=xxx' 'cmd' 'arg0' 'arg1' ...
      command = shlex.split(command)
    else:
      command = [pipes.quote(c) for c in command]
    command = ['ssh'] + options + [remote_sig] + command
    shell = False

    logger.debug('SSHLink: Run [%r]', command)
    close_fds = True
    if (platform.system() == 'Windows' and
        (stdin, stdout, stderr) != (None, None, None)):
      close_fds = False
    proc = subprocess.Popen(command, shell=shell, close_fds=close_fds,
                            stdin=subprocess.PIPE,
                            stdout=stdout, stderr=stderr)
    self._StartWatcher(proc)
    return proc

  def IsReady(self):
    """See DUTLink.IsReady"""
    try:
      return subprocess.call(['ping', '-w', '1', '-c', '1', self.host]) == 0
    except ClientNotExistError:
      return False

  class ControlMasterWatcher(object):
    __metaclass__ = type_utils.Singleton

    def __init__(self, link_instance):
      assert isinstance(link_instance, SSHLink)

      self._link = link_instance
      self._thread = threading.Thread(target=self.Run)
      self._proc_queue = queue.Queue()

      self._user = self._link.user
      self._host = self._link.host
      self._port = self._link.port
      self._link_class_name = self._link.__class__.__name__

    def IsRunning(self):
      if not self._thread:
        return False
      if not self._thread.is_alive():
        self._thread = None
        return False
      return True

    def Start(self):
      if self.IsRunning():
        return

      if self._link.control_persist is None:
        logger.debug('%s %s@%s:%s is not using control master, don\'t start',
                     self._link_class_name, self._user, self._host, self._port)
        return

      self._thread = process_utils.StartDaemonThread(target=self.Run)

    def AddProcess(self, pid, ppid=None):
      """Add an SSH process to monitor.

      If any of added SSH process is still running, ControlMasterWatcher will
      keep monitoring network connectivity.  If network is down, control master
      will be killed.

      Args:
        pid: PID of process using SSH
        ppid: parent PID of given process
      """
      if not self.IsRunning():
        logger.warning('Watcher is not running, %d is not added', pid)
        return
      self._proc_queue.put((pid, ppid))

    def Run(self):
      logger.info('start monitoring control master')

      # an alias to prevent duplicated pylint warnings
      # pylint: disable=protected-access
      _GetLinkSignature = self._link._Signature

      def _IsControlMasterRunning():
        sig, options = _GetLinkSignature(False)
        return subprocess.call(
            ['ssh', '-O', 'check'] + options + [sig, 'true'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE) == 0

      def _StopControlMaster():
        sig, options = _GetLinkSignature(False)
        subprocess.call(['ssh', '-O', 'exit'] + options + [sig, 'true'])

      def _CallTrue():
        sig, options = _GetLinkSignature(False)
        proc = subprocess.Popen(['ssh'] + options + [sig, 'true'])
        time.sleep(1)
        returncode = proc.poll()
        if returncode != 0:
          proc.kill()
          return False
        else:
          return True

      def _PollingCallback(is_process_alive):
        if not is_process_alive:
          return True  # returns True to stop polling

        try:
          if not _IsControlMasterRunning():
            logger.info('control master is not running, skipped')
            return

          if not _CallTrue():
            logger.info('loss connection, stopping control master')
            _StopControlMaster()
        except Exception:
          logger.info('monitoring %s to %s@%s:%s',
                      self._link_class_name,
                      self._user, self._host, self._port, exc_info=True)

      while True:
        # get a new process from queue to monitor
        # since Queue.get will block if queue is empty, we don't need to sleep
        pid, ppid = self._proc_queue.get()
        logger.debug('start monitoring control master until %d terminates', pid)

        sync_utils.PollForCondition(
            lambda: process_utils.IsProcessAlive(pid, ppid),
            condition_method=_PollingCallback,
            timeout_secs=None,
            poll_interval_secs=1)