Add lots of utils, PRESUBMIT.py

py/utils/ssh_utils.py

+#!/usr/bin/env python
+# Copyright (c) 2012 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+""" This module contains tools related to ssh used by the buildbot scripts. """
+
+import atexit
+import os
+import re
+import shell_utils
+import signal
+
+def PutSCP(local_path, remote_path, username, host, port, recurse=False,
+           options=None):
+  """ Send a file to the given host over SCP. Assumes that public key
+  authentication is set up between the client and server.
+
+  local_path: path to the file to send on the client
+  remote_path: destination path for the file on the server
+  username: ssh login name
+  host: hostname or ip address of the server
+  port: port on the server to use
+  recurse: boolean indicating whether to transmit everything in a folder
+  options: list of extra options to pass to scp
+  """
+  # TODO(borenet): This will hang for a while if the host does not recognize
+  # the client.
+  cmd = ['scp']
+  if options:
+    cmd.extend(options)
+  if recurse:
+    cmd.append('-r')
+  cmd.extend(
+    ['-P', port, local_path, '%s@%s:%s' % (username, host, remote_path)])
+  shell_utils.run(cmd)
+
+
+def MultiPutSCP(local_paths, remote_path, username, host, port, options=None):
+  """ Send files to the given host over SCP. Assumes that public key
+  authentication is set up between the client and server.
+
+  local_paths: list of paths of files and directories to send on the client
+  remote_path: destination directory path on the server
+  username: ssh login name
+  host: hostname or ip address of the server
+  port: port on the server to use
+  options: list of extra options to pass to scp
+  """
+  # TODO(borenet): This will hang for a while if the host does not recognize
+  # the client.
+  cmd = ['scp']
+  if options:
+    cmd.extend(options)
+  cmd.extend(['-r', '-P', port])
+  cmd.extend(local_paths)
+  cmd.append('%s@%s:%s' % (username, host, remote_path))
+  shell_utils.run(cmd)
+
+
+def GetSCP(local_path, remote_path, username, host, port, recurse=False,
+           options=None):
+  """ Retrieve a file from the given host over SCP. Assumes that public key
+  authentication is set up between the client and server.
+
+  local_path: destination path for the file on the client
+  remote_path: path to the file to retrieve on the server
+  username: ssh login name
+  host: hostname or ip address of the server
+  port: port on the server to use
+  recurse: boolean indicating whether to transmit everything in a folder
+  options: list of extra options to pass to scp
+  """
+  # TODO(borenet): This will hang for a while if the host does not recognize
+  # the client.
+  cmd = ['scp']
+  if options:
+    cmd.extend(options)
+  if recurse:
+    cmd.append('-r')
+  cmd.extend(
+    ['-P', port, '%s@%s:%s' % (username, host, remote_path), local_path])
+  shell_utils.run(cmd)
+
+
+def RunSSHCmd(username, host, port, command, echo=True, options=None):
+  """ Login to the given host and run the given command.
+
+  username: ssh login name
+  host: hostname or ip address of the server
+  port: port on the server to use
+  command: (string) command to run on the server
+  options: list of extra options to pass to ssh
+  """
+  # TODO(borenet): This will hang for a while if the host does not recognize
+  # the client.
+  cmd = ['ssh']
+  if options:
+    cmd.extend(options)
+  cmd.extend(['-p', port, '%s@%s' % (username, host), command])
+  return shell_utils.run(cmd, echo=echo)
+
+
+def ShellEscape(arg):
+  """ Escape a single argument for passing into a remote shell
+  """
+  arg = re.sub(r'(["\\])', r'\\\1', arg)
+  return '"%s"' % arg if re.search(r'[\' \t\r\n]', arg) else arg
+
+
+def RunSSH(username, host, port, command, echo=True, options=None):
+  """ Login to the given host and run the given command.
+
+  username: ssh login name
+  host: hostname or ip address of the server
+  port: port on the server to use
+  command: command to run on the server in list format
+  options: list of extra options to pass to ssh
+  """
+  cmd = ' '.join(ShellEscape(arg) for arg in command)
+  return RunSSHCmd(username, host, port, cmd, echo=echo, options=options)
+
+
+class SshDestination(object):
+  """ Convenience class to remember a host, port, and username.
+  Wraps the other functions in this module.
+  """
+  def __init__(self, host, port, username, options=None):
+    """
+    host - (string) hostname of the target
+    port - (string or int) sshd port on the target
+    username - (string) remote username
+    options - (list of strings) extra options to pass to ssh and scp.
+    """
+    self.host = host
+    self.port = str(port)
+    self.user = username
+    self.options = options
+
+  def Put(self, local_path, remote_path, recurse=False):
+    return PutSCP(local_path, remote_path, self.user, self.host,
+                  self.port, recurse=recurse, options=self.options)
+
+  def MultiPut(self, local_paths, remote_path):
+    return MultiPutSCP(local_paths, remote_path, self.user, self.host,
+                       self.port, options=self.options)
+
+  def Get(self, local_path, remote_path, recurse=False):
+    return GetSCP(local_path, remote_path, self.user,
+                  self.host, self.port, recurse=recurse, options=self.options)
+
+  def RunCmd(self, command, echo=True):
+    return RunSSHCmd(self.user, self.host, self.port, command,
+                     echo=echo, options=self.options)
+
+  def Run(self, command, echo=True):
+    return RunSSH(self.user, self.host, self.port, command,
+                  echo=echo, options=self.options)
+
+
+def search_within_string(input_string, pattern):
+  """Search for regular expression in a string.
+
+  input_string: (string) to be searched
+  pattern: (string) to be passed to re.compile, with a symbolic
+           group named 'return'.
+  default: what to return if no match
+
+  Returns a string or None
+  """
+  match = re.search(pattern, input_string)
+  return match.group('return') if match else None
+
+def SSHAdd(key_file):
+  """ Call ssh-add, and call ssh-agent if necessary.
+  """
+  assert os.path.isfile(key_file)
+  try:
+    shell_utils.run(['ssh-add', key_file],
+                    log_in_real_time=False)
+    return
+  except shell_utils.CommandFailedException:
+    ssh_agent_output = shell_utils.run(['ssh-agent', '-s'],
+                                       log_in_real_time=False)
+    if not ssh_agent_output:
+      raise Exception('ssh-agent did not print anything')
+    ssh_auth_sock = search_within_string(
+      ssh_agent_output, r'SSH_AUTH_SOCK=(?P<return>[^;]*);')
+    ssh_agent_pid = search_within_string(
+      ssh_agent_output, r'SSH_AGENT_PID=(?P<return>[^;]*);')
+    if not (ssh_auth_sock and ssh_agent_pid):
+      raise Exception('ssh-agent did not print meaningful data')
+    os.environ['SSH_AUTH_SOCK'] = ssh_auth_sock
+    os.environ['SSH_AGENT_PID'] = ssh_agent_pid
+    atexit.register(os.kill, int(ssh_agent_pid), signal.SIGTERM)
+    shell_utils.run(['ssh-add', key_file])