• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2011 Google Inc. All Rights Reserved.
2
3__author__ = 'kbaclawski@google.com (Krystian Baclawski)'
4
5import abc
6import collections
7import os.path
8
9
10class 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
36class 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
74class 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
127class 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
138class 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
172def 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
188def 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
202def MakeSymlink(to_path, link_name):
203  return Shell('ln', '-f', '-s', '-T', to_path, link_name)
204
205
206def 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
219def RmTree(*dirs):
220  return Shell('rm', '-r', '-f', *dirs)
221
222
223def UnTar(tar_file, dest_dir):
224  return Chain(
225      MakeDir(dest_dir), Shell('tar', '-x', '-f', tar_file, '-C', dest_dir))
226
227
228def 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