automation/common/command.py

# Copyright 2011 Google Inc. All Rights Reserved.

__author__ = 'kbaclawski@google.com (Krystian Baclawski)'

import abc
import collections
import os.path


class Shell(object):
  """Class used to build a string representation of a shell command."""

  def __init__(self, cmd, *args, **kwargs):
    assert all(key in ['path', 'ignore_error'] for key in kwargs)

    self._cmd = cmd
    self._args = list(args)
    self._path = kwargs.get('path', '')
    self._ignore_error = bool(kwargs.get('ignore_error', False))

  def __str__(self):
    cmdline = [os.path.join(self._path, self._cmd)]
    cmdline.extend(self._args)

    cmd = ' '.join(cmdline)

    if self._ignore_error:
      cmd = '{ %s; true; }' % cmd

    return cmd

  def AddOption(self, option):
    self._args.append(option)


class Wrapper(object):
  """Wraps a command with environment which gets cleaned up after execution."""

  _counter = 1

  def __init__(self, command, cwd=None, env=None, umask=None):
    # @param cwd: temporary working directory
    # @param env: dictionary of environment variables
    self._command = command
    self._prefix = Chain()
    self._suffix = Chain()

    if cwd:
      self._prefix.append(Shell('pushd', cwd))
      self._suffix.insert(0, Shell('popd'))

    if env:
      for env_var, value in env.items():
        self._prefix.append(Shell('%s=%s' % (env_var, value)))
        self._suffix.insert(0, Shell('unset', env_var))

    if umask:
      umask_save_var = 'OLD_UMASK_%d' % self.counter

      self._prefix.append(Shell('%s=$(umask)' % umask_save_var))
      self._prefix.append(Shell('umask', umask))
      self._suffix.insert(0, Shell('umask', '$%s' % umask_save_var))

  @property
  def counter(self):
    counter = self._counter
    self._counter += 1
    return counter

  def __str__(self):
    return str(Chain(self._prefix, self._command, self._suffix))


class AbstractCommandContainer(collections.MutableSequence):
  """Common base for all classes that behave like command container."""

  def __init__(self, *commands):
    self._commands = list(commands)

  def __contains__(self, command):
    return command in self._commands

  def __iter__(self):
    return iter(self._commands)

  def __len__(self):
    return len(self._commands)

  def __getitem__(self, index):
    return self._commands[index]

  def __setitem__(self, index, command):
    self._commands[index] = self._ValidateCommandType(command)

  def __delitem__(self, index):
    del self._commands[index]

  def insert(self, index, command):
    self._commands.insert(index, self._ValidateCommandType(command))

  @abc.abstractmethod
  def __str__(self):
    pass

  @abc.abstractproperty
  def stored_types(self):
    pass

  def _ValidateCommandType(self, command):
    if type(command) not in self.stored_types:
      raise TypeError('Command cannot have %s type.' % type(command))
    else:
      return command

  def _StringifyCommands(self):
    cmds = []

    for cmd in self:
      if isinstance(cmd, AbstractCommandContainer) and len(cmd) > 1:
        cmds.append('{ %s; }' % cmd)
      else:
        cmds.append(str(cmd))

    return cmds


class Chain(AbstractCommandContainer):
  """Container that chains shell commands using (&&) shell operator."""

  @property
  def stored_types(self):
    return [str, Shell, Chain, Pipe]

  def __str__(self):
    return ' && '.join(self._StringifyCommands())


class Pipe(AbstractCommandContainer):
  """Container that chains shell commands using pipe (|) operator."""

  def __init__(self, *commands, **kwargs):
    assert all(key in ['input', 'output'] for key in kwargs)

    AbstractCommandContainer.__init__(self, *commands)

    self._input = kwargs.get('input', None)
    self._output = kwargs.get('output', None)

  @property
  def stored_types(self):
    return [str, Shell]

  def __str__(self):
    pipe = self._StringifyCommands()

    if self._input:
      pipe.insert(str(Shell('cat', self._input), 0))

    if self._output:
      pipe.append(str(Shell('tee', self._output)))

    return ' | '.join(pipe)

# TODO(kbaclawski): Unfortunately we don't have any policy describing which
# directories can or cannot be touched by a job. Thus, I cannot decide how to
# protect a system against commands that are considered to be dangerous (like
# RmTree("${HOME}")). AFAIK we'll have to execute some commands with root access
# (especially for ChromeOS related jobs, which involve chroot-ing), which is
# even more scary.


def Copy(*args, **kwargs):
  assert all(key in ['to_dir', 'recursive'] for key in kwargs.keys())

  options = []

  if 'to_dir' in kwargs:
    options.extend(['-t', kwargs['to_dir']])

  if 'recursive' in kwargs:
    options.append('-r')

  options.extend(args)

  return Shell('cp', *options)


def RemoteCopyFrom(from_machine, from_path, to_path, username=None):
  from_path = os.path.expanduser(from_path) + '/'
  to_path = os.path.expanduser(to_path) + '/'

  if not username:
    login = from_machine
  else:
    login = '%s@%s' % (username, from_machine)

  return Chain(
      MakeDir(to_path), Shell('rsync', '-a', '%s:%s' %
                              (login, from_path), to_path))


def MakeSymlink(to_path, link_name):
  return Shell('ln', '-f', '-s', '-T', to_path, link_name)


def MakeDir(*dirs, **kwargs):
  options = ['-p']

  mode = kwargs.get('mode', None)

  if mode:
    options.extend(['-m', str(mode)])

  options.extend(dirs)

  return Shell('mkdir', *options)


def RmTree(*dirs):
  return Shell('rm', '-r', '-f', *dirs)


def UnTar(tar_file, dest_dir):
  return Chain(
      MakeDir(dest_dir), Shell('tar', '-x', '-f', tar_file, '-C', dest_dir))


def Tar(tar_file, *args):
  options = ['-c']

  if tar_file.endswith('.tar.bz2'):
    options.append('-j')
  elif tar_file.endswith('.tar.gz'):
    options.append('-z')
  else:
    assert tar_file.endswith('.tar')

  options.extend(['-f', tar_file])
  options.extend(args)

  return Chain(MakeDir(os.path.dirname(tar_file)), Shell('tar', *options))