Home | History | Annotate | Download | only in common 
      1 # Copyright 2011 Google Inc. All Rights Reserved.
      2 
      3 __author__ = 'kbaclawski (at] google.com (Krystian Baclawski)'
      4 
      5 import abc
      6 import collections
      7 import os.path
      8 
      9 
     10 class Shell(object):
     11   """Class used to build a string representation of a shell command."""
     12 
     13   def __init__(self, cmd, *args, **kwargs):
     14     assert all(key in ['path', 'ignore_error'] for key in kwargs)
     15 
     16     self._cmd = cmd
     17     self._args = list(args)
     18     self._path = kwargs.get('path', '')
     19     self._ignore_error = bool(kwargs.get('ignore_error', False))
     20 
     21   def __str__(self):
     22     cmdline = [os.path.join(self._path, self._cmd)]
     23     cmdline.extend(self._args)
     24 
     25     cmd = ' '.join(cmdline)
     26 
     27     if self._ignore_error:
     28       cmd = '{ %s; true; }' % cmd
     29 
     30     return cmd
     31 
     32   def AddOption(self, option):
     33     self._args.append(option)
     34 
     35 
     36 class Wrapper(object):
     37   """Wraps a command with environment which gets cleaned up after execution."""
     38 
     39   _counter = 1
     40 
     41   def __init__(self, command, cwd=None, env=None, umask=None):
     42     # @param cwd: temporary working directory
     43     # @param env: dictionary of environment variables
     44     self._command = command
     45     self._prefix = Chain()
     46     self._suffix = Chain()
     47 
     48     if cwd:
     49       self._prefix.append(Shell('pushd', cwd))
     50       self._suffix.insert(0, Shell('popd'))
     51 
     52     if env:
     53       for env_var, value in env.items():
     54         self._prefix.append(Shell('%s=%s' % (env_var, value)))
     55         self._suffix.insert(0, Shell('unset', env_var))
     56 
     57     if umask:
     58       umask_save_var = 'OLD_UMASK_%d' % self.counter
     59 
     60       self._prefix.append(Shell('%s=$(umask)' % umask_save_var))
     61       self._prefix.append(Shell('umask', umask))
     62       self._suffix.insert(0, Shell('umask', '$%s' % umask_save_var))
     63 
     64   @property
     65   def counter(self):
     66     counter = self._counter
     67     self._counter += 1
     68     return counter
     69 
     70   def __str__(self):
     71     return str(Chain(self._prefix, self._command, self._suffix))
     72 
     73 
     74 class AbstractCommandContainer(collections.MutableSequence):
     75   """Common base for all classes that behave like command container."""
     76 
     77   def __init__(self, *commands):
     78     self._commands = list(commands)
     79 
     80   def __contains__(self, command):
     81     return command in self._commands
     82 
     83   def __iter__(self):
     84     return iter(self._commands)
     85 
     86   def __len__(self):
     87     return len(self._commands)
     88 
     89   def __getitem__(self, index):
     90     return self._commands[index]
     91 
     92   def __setitem__(self, index, command):
     93     self._commands[index] = self._ValidateCommandType(command)
     94 
     95   def __delitem__(self, index):
     96     del self._commands[index]
     97 
     98   def insert(self, index, command):
     99     self._commands.insert(index, self._ValidateCommandType(command))
    100 
    101   @abc.abstractmethod
    102   def __str__(self):
    103     pass
    104 
    105   @abc.abstractproperty
    106   def stored_types(self):
    107     pass
    108 
    109   def _ValidateCommandType(self, command):
    110     if type(command) not in self.stored_types:
    111       raise TypeError('Command cannot have %s type.' % type(command))
    112     else:
    113       return command
    114 
    115   def _StringifyCommands(self):
    116     cmds = []
    117 
    118     for cmd in self:
    119       if isinstance(cmd, AbstractCommandContainer) and len(cmd) > 1:
    120         cmds.append('{ %s; }' % cmd)
    121       else:
    122         cmds.append(str(cmd))
    123 
    124     return cmds
    125 
    126 
    127 class Chain(AbstractCommandContainer):
    128   """Container that chains shell commands using (&&) shell operator."""
    129 
    130   @property
    131   def stored_types(self):
    132     return [str, Shell, Chain, Pipe]
    133 
    134   def __str__(self):
    135     return ' && '.join(self._StringifyCommands())
    136 
    137 
    138 class Pipe(AbstractCommandContainer):
    139   """Container that chains shell commands using pipe (|) operator."""
    140 
    141   def __init__(self, *commands, **kwargs):
    142     assert all(key in ['input', 'output'] for key in kwargs)
    143 
    144     AbstractCommandContainer.__init__(self, *commands)
    145 
    146     self._input = kwargs.get('input', None)
    147     self._output = kwargs.get('output', None)
    148 
    149   @property
    150   def stored_types(self):
    151     return [str, Shell]
    152 
    153   def __str__(self):
    154     pipe = self._StringifyCommands()
    155 
    156     if self._input:
    157       pipe.insert(str(Shell('cat', self._input), 0))
    158 
    159     if self._output:
    160       pipe.append(str(Shell('tee', self._output)))
    161 
    162     return ' | '.join(pipe)
    163 
    164 # TODO(kbaclawski): Unfortunately we don't have any policy describing which
    165 # directories can or cannot be touched by a job. Thus, I cannot decide how to
    166 # protect a system against commands that are considered to be dangerous (like
    167 # RmTree("${HOME}")). AFAIK we'll have to execute some commands with root access
    168 # (especially for ChromeOS related jobs, which involve chroot-ing), which is
    169 # even more scary.
    170 
    171 
    172 def Copy(*args, **kwargs):
    173   assert all(key in ['to_dir', 'recursive'] for key in kwargs.keys())
    174 
    175   options = []
    176 
    177   if 'to_dir' in kwargs:
    178     options.extend(['-t', kwargs['to_dir']])
    179 
    180   if 'recursive' in kwargs:
    181     options.append('-r')
    182 
    183   options.extend(args)
    184 
    185   return Shell('cp', *options)
    186 
    187 
    188 def RemoteCopyFrom(from_machine, from_path, to_path, username=None):
    189   from_path = os.path.expanduser(from_path) + '/'
    190   to_path = os.path.expanduser(to_path) + '/'
    191 
    192   if not username:
    193     login = from_machine
    194   else:
    195     login = '%s@%s' % (username, from_machine)
    196 
    197   return Chain(
    198       MakeDir(to_path), Shell('rsync', '-a', '%s:%s' %
    199                               (login, from_path), to_path))
    200 
    201 
    202 def MakeSymlink(to_path, link_name):
    203   return Shell('ln', '-f', '-s', '-T', to_path, link_name)
    204 
    205 
    206 def MakeDir(*dirs, **kwargs):
    207   options = ['-p']
    208 
    209   mode = kwargs.get('mode', None)
    210 
    211   if mode:
    212     options.extend(['-m', str(mode)])
    213 
    214   options.extend(dirs)
    215 
    216   return Shell('mkdir', *options)
    217 
    218 
    219 def RmTree(*dirs):
    220   return Shell('rm', '-r', '-f', *dirs)
    221 
    222 
    223 def UnTar(tar_file, dest_dir):
    224   return Chain(
    225       MakeDir(dest_dir), Shell('tar', '-x', '-f', tar_file, '-C', dest_dir))
    226 
    227 
    228 def Tar(tar_file, *args):
    229   options = ['-c']
    230 
    231   if tar_file.endswith('.tar.bz2'):
    232     options.append('-j')
    233   elif tar_file.endswith('.tar.gz'):
    234     options.append('-z')
    235   else:
    236     assert tar_file.endswith('.tar')
    237 
    238   options.extend(['-f', tar_file])
    239   options.extend(args)
    240 
    241   return Chain(MakeDir(os.path.dirname(tar_file)), Shell('tar', *options))
    242