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

import argparse
import glob
import os
import re
import shlex
import subprocess

from host import Host


class Device(object):
  """Represents a Fuchsia device attached to a host.

    This class abstracts the details of remotely running commands and
    transferring data to and from the device.

    Attributes:
      host: A Host object represent the local platform attached to this target
        device.
  """

  @classmethod
  def from_args(cls, host, args):
    """Constructs a Device from command line arguments."""
    netaddr_cmd = ['netaddr', '--fuchsia', '--nowait']
    default_device = '{}.device'.format(host.build_dir)
    if args.device:
      netaddr_cmd.append(args.device)
    elif os.path.exists(default_device):
      with open(default_device) as f:
        netaddr_cmd.append(f.read().strip())
    try:
      netaddr = host.zircon_tool(netaddr_cmd)
    except subprocess.CalledProcessError:
      raise RuntimeError('Unable to find device; try `fx set-device`.')
    device = cls(host, netaddr)
    if not host.build_dir:
      raise Host.ConfigError('Unable to find SSH configuration.')
    device.set_ssh_config(Host.join(host.build_dir, 'ssh-keys', 'ssh_config'))
    return device

  def __init__(self, host, addr, port=22):
    self.host = host
    self._addr = addr
    self._ssh_opts = {}
    if port != 22:
      self._ssh_opts['p'] = [str(port)]

  def set_ssh_config(self, config_file):
    """Sets the SSH arguments to use a config file."""
    if not os.path.exists(config_file):
      raise Host.ConfigError('Unable to find SSH configuration.')
    self._ssh_opts['F'] = [config_file]

  def set_ssh_identity(self, identity_file):
    if not os.path.exists(identity_file):
      raise Host.ConfigError('Unable to find SSH identity.')
    self._ssh_opts['i'] = [identity_file]

  def set_ssh_option(self, option):
    """Sets SSH configuration options. Can be used multiple times."""
    if 'o' in self._ssh_opts:
      self._ssh_opts['o'].append(option)
    else:
      self._ssh_opts['o'] = [option]

  def set_ssh_verbosity(self, level):
    """Sets how much debugging SSH prints. Default is 0 (none), max is 3."""
    for i in range(1, 4):
      opt = 'v' * i
      if level == i and not opt in self._ssh_opts:
        self._ssh_opts[opt] = []
      elif level != i and opt in self._ssh_opts:
        del self._ssh_opts[opt]

  def get_ssh_cmd(self, cmd):
    """Returns the SSH executable and options."""
    result = cmd[:1]
    for opt, args in self._ssh_opts.iteritems():
      if len(args) == 0:
        result.append('-' + opt)
      else:
        for arg in args:
          result.append('-' + opt)
          result.append(arg)
    return result + cmd[1:]

  def _ssh(self, cmdline, stdout=subprocess.PIPE):
    """Internal wrapper around _rexec that adds the ssh command and config.

    Don't call this directly. This method exists to be overridden in testing.

    Args:
      cmdline: List of command line arguments to execute on device
      stdout: Same as for subprocess.Popen

    Returns:
      If check was false, a subprocess.Popen object representing the running
      child process.

    Raises: Same as subprocess.Popen
    """
    return subprocess.Popen(
        self.get_ssh_cmd(['ssh', self._addr] + cmdline),
        stdout=stdout,
        stderr=subprocess.STDOUT)

  def ssh(self, cmdline, quiet=True, logfile=None):
    """Runs a command to completion on the device.

    Connects to the target device and executes a shell command.  Output from
    the shell command is sent to stdout, and may optionally be saved to a file
    via the POSIX utility 'tee'.

    Args:
      cmdline: A list of command line arguments, starting with the command to
        execute.
      logfile: An optional pathname to save a copy of the command output to. The
        output will also still be sent to stdout.
    """
    if quiet:
      if logfile:
        with open(logfile, 'w') as f:
          self._ssh(cmdline, stdout=f).wait()
      else:
        self._ssh(cmdline, stdout=Host.DEVNULL).wait()
    else:
      if logfile:
        proc = self._ssh(cmdline, stdout=subprocess.PIPE)
        subprocess.check_call(['tee', logfile], stdin=proc.stdout)
      else:
        self._ssh(cmdline, stdout=None).wait()

  def getpids(self):
    """Maps names to process IDs for running fuzzers.

    Connects to the device and checks which fuzz targets have a matching entry
    in the component list given by 'cs'.  This matches on *only* the first 32
    characters of the component manifest and package URL.  This is due to 'cs'
    being limited to returning strings of length `ZX_MAX_NAME_LEN`, as defined
    in //zircon/system/public/zircon/types.h.

    Returns:
      A dict mapping fuzz target names to process IDs. May be empty if no
      fuzzers are running.
    """
    out, _ = self._ssh(['cs'], stdout=subprocess.PIPE).communicate()
    pids = {}
    for fuzzer in self.host.fuzzers:
      tgt = (fuzzer[1] + '.cmx')[:32]
      url = ('fuchsia-pkg://fuchsia.com/%s#meta' % fuzzer[0])[:32]
      for line in str(out).split('\n'):
        match = re.search(tgt + r'\[(\d+)\]: ' + url, line)
        if match:
          pids[fuzzer[1]] = int(match.group(1))
    return pids

  def ls(self, path):
    """Maps file names to sizes for the given path.

    Connects to a Fuchsia device and lists the files in a directory given by
    the provided path.  Ignore non-existent paths.

    Args:
      path: Absolute path to a directory on the device.

    Returns:
      A dict mapping file names to file sizes, or an empty dict if the path
      does not exist.
    """
    results = {}
    try:
      out, _ = self._ssh(
          ['ls', '-l', path], stdout=subprocess.PIPE).communicate()
      for line in str(out).split('\n'):
        # Line ~= '-rw-r--r-- 1 0 0 8192 Mar 18 22:02 some-name'
        parts = line.split()
        if len(parts) > 8:
          results[' '.join(parts[8:])] = int(parts[4])
    except subprocess.CalledProcessError:
      pass
    return results

  def _scp(self, srcs, dst):
    """Copies `src` to `dst`.

    Don't call directly; use `fetch` or `store` instead.`

    Args:
      srcs: Local or remote paths to copy from.
      dst: Local or remote path to copy to.
    """
    # Wild cards only work if shell=True and the whole line passed as a string.
    cmd = self.get_ssh_cmd(['scp'] + srcs + [dst])
    subprocess.check_call(cmd, stdout=None, stderr=None)

  def fetch(self, data_src, host_dst):
    """Copies `data_src` on the target to `host_dst` on the host."""
    if not os.path.isdir(host_dst):
      raise ValueError(host_dst + ' is not a directory')
    self._scp(['[{}]:{}'.format(self._addr, data_src)], host_dst)

  def store(self, host_src, data_dst):
    """Copies `host_src` on the host to `data_dst` on the target."""
    self.ssh(['mkdir', '-p', data_dst])
    self._scp(glob.glob(host_src), '[{}]:{}'.format(self._addr, data_dst))