command_executer.py revision a8af9a7a2462b00e72deff99327bdb452a715277 
1# Copyright 2011 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4"""Utilities to run commands in outside/inside chroot and on the board."""
5
6from __future__ import print_function
7
8import getpass
9import os
10import re
11import select
12import signal
13import subprocess
14import sys
15import tempfile
16import time
17
18import logger
19import misc
20
21mock_default = False
22
23LOG_LEVEL = ('none', 'quiet', 'average', 'verbose')
24
25
26def InitCommandExecuter(mock=False):
27  # pylint: disable=global-statement
28  global mock_default
29  # Whether to default to a mock command executer or not
30  mock_default = mock
31
32
33def GetCommandExecuter(logger_to_set=None, mock=False, log_level='verbose'):
34  # If the default is a mock executer, always return one.
35  if mock_default or mock:
36    return MockCommandExecuter(log_level, logger_to_set)
37  else:
38    return CommandExecuter(log_level, logger_to_set)
39
40
41class CommandExecuter(object):
42  """Provides several methods to execute commands on several environments."""
43
44  def __init__(self, log_level, logger_to_set=None):
45    self.log_level = log_level
46    if log_level == 'none':
47      self.logger = None
48    else:
49      if logger_to_set is not None:
50        self.logger = logger_to_set
51      else:
52        self.logger = logger.GetLogger()
53
54  def GetLogLevel(self):
55    return self.log_level
56
57  def SetLogLevel(self, log_level):
58    self.log_level = log_level
59
60  def RunCommandGeneric(self,
61                        cmd,
62                        return_output=False,
63                        machine=None,
64                        username=None,
65                        command_terminator=None,
66                        command_timeout=None,
67                        terminated_timeout=10,
68                        print_to_console=True,
69                        except_handler=lambda p, e: None):
70    """Run a command.
71
72    Returns triplet (returncode, stdout, stderr).
73    """
74
75    cmd = str(cmd)
76
77    if self.log_level == 'quiet':
78      print_to_console = False
79
80    if self.log_level == 'verbose':
81      self.logger.LogCmd(cmd, machine, username, print_to_console)
82    elif self.logger:
83      self.logger.LogCmdToFileOnly(cmd, machine, username)
84    if command_terminator and command_terminator.IsTerminated():
85      if self.logger:
86        self.logger.LogError('Command was terminated!', print_to_console)
87      return (1, '', '')
88
89    if machine is not None:
90      user = ''
91      if username is not None:
92        user = username + '@'
93      cmd = "ssh -t -t %s%s -- '%s'" % (user, machine, cmd)
94
95    # We use setsid so that the child will have a different session id
96    # and we can easily kill the process group. This is also important
97    # because the child will be disassociated from the parent terminal.
98    # In this way the child cannot mess the parent's terminal.
99    p = None
100    try:
101      p = subprocess.Popen(cmd,
102                           stdout=subprocess.PIPE,
103                           stderr=subprocess.PIPE,
104                           shell=True,
105                           preexec_fn=os.setsid)
106
107      full_stdout = ''
108      full_stderr = ''
109
110      # Pull output from pipes, send it to file/stdout/string
111      out = err = None
112      pipes = [p.stdout, p.stderr]
113
114      my_poll = select.poll()
115      my_poll.register(p.stdout, select.POLLIN)
116      my_poll.register(p.stderr, select.POLLIN)
117
118      terminated_time = None
119      started_time = time.time()
120
121      while len(pipes):
122        if command_terminator and command_terminator.IsTerminated():
123          os.killpg(os.getpgid(p.pid), signal.SIGTERM)
124          if self.logger:
125            self.logger.LogError('Command received termination request. '
126                                 'Killed child process group.',
127                                 print_to_console)
128          break
129
130        l = my_poll.poll(100)
131        for (fd, _) in l:
132          if fd == p.stdout.fileno():
133            out = os.read(p.stdout.fileno(), 16384)
134            if return_output:
135              full_stdout += out
136            if self.logger:
137              self.logger.LogCommandOutput(out, print_to_console)
138            if out == '':
139              pipes.remove(p.stdout)
140              my_poll.unregister(p.stdout)
141          if fd == p.stderr.fileno():
142            err = os.read(p.stderr.fileno(), 16384)
143            if return_output:
144              full_stderr += err
145            if self.logger:
146              self.logger.LogCommandError(err, print_to_console)
147            if err == '':
148              pipes.remove(p.stderr)
149              my_poll.unregister(p.stderr)
150
151        if p.poll() is not None:
152          if terminated_time is None:
153            terminated_time = time.time()
154          elif (terminated_timeout is not None and
155                time.time() - terminated_time > terminated_timeout):
156            if self.logger:
157              self.logger.LogWarning('Timeout of %s seconds reached since '
158                                     'process termination.' %
159                                     terminated_timeout,
160                                     print_to_console)
161            break
162
163        if (command_timeout is not None and
164            time.time() - started_time > command_timeout):
165          os.killpg(os.getpgid(p.pid), signal.SIGTERM)
166          if self.logger:
167            self.logger.LogWarning('Timeout of %s seconds reached since process'
168                                   'started. Killed child process group.' %
169                                   command_timeout, print_to_console)
170          break
171
172        if out == err == '':
173          break
174
175      p.wait()
176      if return_output:
177        return (p.returncode, full_stdout, full_stderr)
178      return (p.returncode, '', '')
179    except BaseException as e:
180      except_handler(p, e)
181      raise
182
183  def RunCommand(self, *args, **kwargs):
184    """Run a command.
185
186    Takes the same arguments as RunCommandGeneric except for return_output.
187    Returns a single value returncode.
188    """
189    # Make sure that args does not overwrite 'return_output'
190    assert len(args) <= 1
191    assert 'return_output' not in kwargs
192    kwargs['return_output'] = False
193    return self.RunCommandGeneric(*args, **kwargs)[0]
194
195  def RunCommandWExceptionCleanup(self, *args, **kwargs):
196    """Run a command and kill process if exception is thrown.
197
198    Takes the same arguments as RunCommandGeneric except for except_handler.
199    Returns same as RunCommandGeneric.
200    """
201
202    def KillProc(proc, _):
203      if proc:
204        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
205
206    # Make sure that args does not overwrite 'except_handler'
207    assert len(args) <= 8
208    assert 'except_handler' not in kwargs
209    kwargs['except_handler'] = KillProc
210    return self.RunCommandGeneric(*args, **kwargs)
211
212  def RunCommandWOutput(self, *args, **kwargs):
213    """Run a command.
214
215    Takes the same arguments as RunCommandGeneric except for return_output.
216    Returns a triplet (returncode, stdout, stderr).
217    """
218    # Make sure that args does not overwrite 'return_output'
219    assert len(args) <= 1
220    assert 'return_output' not in kwargs
221    kwargs['return_output'] = True
222    return self.RunCommandGeneric(*args, **kwargs)
223
224  def RemoteAccessInitCommand(self, chromeos_root, machine):
225    command = ''
226    command += '\nset -- --remote=' + machine
227    command += '\n. ' + chromeos_root + '/src/scripts/common.sh'
228    command += '\n. ' + chromeos_root + '/src/scripts/remote_access.sh'
229    command += '\nTMP=$(mktemp -d)'
230    command += "\nFLAGS \"$@\" || exit 1"
231    command += '\nremote_access_init'
232    return command
233
234  def WriteToTempShFile(self, contents):
235    handle, command_file = tempfile.mkstemp(prefix=os.uname()[1], suffix='.sh')
236    os.write(handle, '#!/bin/bash\n')
237    os.write(handle, contents)
238    os.close(handle)
239    return command_file
240
241  def CrosLearnBoard(self, chromeos_root, machine):
242    command = self.RemoteAccessInitCommand(chromeos_root, machine)
243    command += '\nlearn_board'
244    command += '\necho ${FLAGS_board}'
245    retval, output, _ = self.RunCommandWOutput(command)
246    if self.logger:
247      self.logger.LogFatalIf(retval, 'learn_board command failed')
248    elif retval:
249      sys.exit(1)
250    return output.split()[-1]
251
252  def CrosRunCommandGeneric(self,
253                            cmd,
254                            return_output=False,
255                            machine=None,
256                            command_terminator=None,
257                            chromeos_root=None,
258                            command_timeout=None,
259                            terminated_timeout=10,
260                            print_to_console=True):
261    """Run a command on a ChromeOS box.
262
263    Returns triplet (returncode, stdout, stderr).
264    """
265
266    if self.log_level != 'verbose':
267      print_to_console = False
268
269    if self.logger:
270      self.logger.LogCmd(cmd, print_to_console=print_to_console)
271      self.logger.LogFatalIf(not machine, 'No machine provided!')
272      self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!')
273    else:
274      if not chromeos_root or not machine:
275        sys.exit(1)
276    chromeos_root = os.path.expanduser(chromeos_root)
277
278    # Write all commands to a file.
279    command_file = self.WriteToTempShFile(cmd)
280    retval = self.CopyFiles(command_file,
281                            command_file,
282                            dest_machine=machine,
283                            command_terminator=command_terminator,
284                            chromeos_root=chromeos_root,
285                            dest_cros=True,
286                            recursive=False,
287                            print_to_console=print_to_console)
288    if retval:
289      if self.logger:
290        self.logger.LogError('Could not run remote command on machine.'
291                             ' Is the machine up?')
292      return (retval, '', '')
293
294    command = self.RemoteAccessInitCommand(chromeos_root, machine)
295    command += '\nremote_sh bash %s' % command_file
296    command += "\nl_retval=$?; echo \"$REMOTE_OUT\"; exit $l_retval"
297    retval = self.RunCommandGeneric(command,
298                                    return_output,
299                                    command_terminator=command_terminator,
300                                    command_timeout=command_timeout,
301                                    terminated_timeout=terminated_timeout,
302                                    print_to_console=print_to_console)
303    if return_output:
304      connect_signature = (
305          'Initiating first contact with remote host\n' + 'Connection OK\n')
306      connect_signature_re = re.compile(connect_signature)
307      modded_retval = list(retval)
308      modded_retval[1] = connect_signature_re.sub('', retval[1])
309      return modded_retval
310    return retval
311
312  def CrosRunCommand(self, *args, **kwargs):
313    """Run a command on a ChromeOS box.
314
315    Takes the same arguments as CrosRunCommandGeneric except for return_output.
316    Returns a single value returncode.
317    """
318    # Make sure that args does not overwrite 'return_output'
319    assert len(args) <= 1
320    assert 'return_output' not in kwargs
321    kwargs['return_output'] = False
322    return self.CrosRunCommandGeneric(*args, **kwargs)[0]
323
324  def CrosRunCommandWOutput(self, *args, **kwargs):
325    """Run a command on a ChromeOS box.
326
327    Takes the same arguments as CrosRunCommandGeneric except for return_output.
328    Returns a triplet (returncode, stdout, stderr).
329    """
330    # Make sure that args does not overwrite 'return_output'
331    assert len(args) <= 1
332    assert 'return_output' not in kwargs
333    kwargs['return_output'] = True
334    return self.CrosRunCommandGeneric(*args, **kwargs)
335
336  def ChrootRunCommandGeneric(self,
337                              chromeos_root,
338                              command,
339                              return_output=False,
340                              command_terminator=None,
341                              command_timeout=None,
342                              terminated_timeout=10,
343                              print_to_console=True,
344                              cros_sdk_options=''):
345    """Runs a command within the chroot.
346
347    Returns triplet (returncode, stdout, stderr).
348    """
349
350    if self.log_level != 'verbose':
351      print_to_console = False
352
353    if self.logger:
354      self.logger.LogCmd(command, print_to_console=print_to_console)
355
356    handle, command_file = tempfile.mkstemp(
357        dir=os.path.join(chromeos_root, 'src/scripts'),
358        suffix='.sh',
359        prefix='in_chroot_cmd')
360    os.write(handle, '#!/bin/bash\n')
361    os.write(handle, command)
362    os.write(handle, '\n')
363    os.close(handle)
364
365    os.chmod(command_file, 0777)
366
367    # if return_output is set, run a dummy command first to make sure that
368    # the chroot already exists. We want the final returned output to skip
369    # the output from chroot creation steps.
370    if return_output:
371      ret = self.RunCommand('cd %s; cros_sdk %s -- true' %
372                            (chromeos_root, cros_sdk_options))
373      if ret:
374        return (ret, '', '')
375
376    # Run command_file inside the chroot, making sure that any "~" is expanded
377    # by the shell inside the chroot, not outside.
378    command = ("cd %s; cros_sdk %s -- bash -c '%s/%s'" %
379               (chromeos_root, cros_sdk_options, misc.CHROMEOS_SCRIPTS_DIR,
380                os.path.basename(command_file)))
381    ret = self.RunCommandGeneric(command,
382                                 return_output,
383                                 command_terminator=command_terminator,
384                                 command_timeout=command_timeout,
385                                 terminated_timeout=terminated_timeout,
386                                 print_to_console=print_to_console)
387    os.remove(command_file)
388    return ret
389
390  def ChrootRunCommand(self, *args, **kwargs):
391    """Runs a command within the chroot.
392
393    Takes the same arguments as ChrootRunCommandGeneric except for
394    return_output.
395    Returns a single value returncode.
396    """
397    # Make sure that args does not overwrite 'return_output'
398    assert len(args) <= 2
399    assert 'return_output' not in kwargs
400    kwargs['return_output'] = False
401    return self.ChrootRunCommandGeneric(*args, **kwargs)[0]
402
403  def ChrootRunCommandWOutput(self, *args, **kwargs):
404    """Runs a command within the chroot.
405
406    Takes the same arguments as ChrootRunCommandGeneric except for
407    return_output.
408    Returns a triplet (returncode, stdout, stderr).
409    """
410    # Make sure that args does not overwrite 'return_output'
411    assert len(args) <= 2
412    assert 'return_output' not in kwargs
413    kwargs['return_output'] = True
414    return self.ChrootRunCommandGeneric(*args, **kwargs)
415
416  def RunCommands(self,
417                  cmdlist,
418                  machine=None,
419                  username=None,
420                  command_terminator=None):
421    cmd = ' ;\n'.join(cmdlist)
422    return self.RunCommand(cmd,
423                           machine=machine,
424                           username=username,
425                           command_terminator=command_terminator)
426
427  def CopyFiles(self,
428                src,
429                dest,
430                src_machine=None,
431                dest_machine=None,
432                src_user=None,
433                dest_user=None,
434                recursive=True,
435                command_terminator=None,
436                chromeos_root=None,
437                src_cros=False,
438                dest_cros=False,
439                print_to_console=True):
440    src = os.path.expanduser(src)
441    dest = os.path.expanduser(dest)
442
443    if recursive:
444      src = src + '/'
445      dest = dest + '/'
446
447    if src_cros == True or dest_cros == True:
448      if self.logger:
449        self.logger.LogFatalIf(src_cros == dest_cros,
450                               'Only one of src_cros and desc_cros can '
451                               'be True.')
452        self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!')
453      elif src_cros == dest_cros or not chromeos_root:
454        sys.exit(1)
455      if src_cros == True:
456        cros_machine = src_machine
457      else:
458        cros_machine = dest_machine
459
460      command = self.RemoteAccessInitCommand(chromeos_root, cros_machine)
461      ssh_command = (
462          'ssh -p ${FLAGS_ssh_port}' + ' -o StrictHostKeyChecking=no' +
463          ' -o UserKnownHostsFile=$(mktemp)' + ' -i $TMP_PRIVATE_KEY')
464      rsync_prefix = "\nrsync -r -e \"%s\" " % ssh_command
465      if dest_cros == True:
466        command += rsync_prefix + '%s root@%s:%s' % (src, dest_machine, dest)
467        return self.RunCommand(command,
468                               machine=src_machine,
469                               username=src_user,
470                               command_terminator=command_terminator,
471                               print_to_console=print_to_console)
472      else:
473        command += rsync_prefix + 'root@%s:%s %s' % (src_machine, src, dest)
474        return self.RunCommand(command,
475                               machine=dest_machine,
476                               username=dest_user,
477                               command_terminator=command_terminator,
478                               print_to_console=print_to_console)
479
480    if dest_machine == src_machine:
481      command = 'rsync -a %s %s' % (src, dest)
482    else:
483      if src_machine is None:
484        src_machine = os.uname()[1]
485        src_user = getpass.getuser()
486      command = 'rsync -a %s@%s:%s %s' % (src_user, src_machine, src, dest)
487    return self.RunCommand(command,
488                           machine=dest_machine,
489                           username=dest_user,
490                           command_terminator=command_terminator,
491                           print_to_console=print_to_console)
492
493  def RunCommand2(self,
494                  cmd,
495                  cwd=None,
496                  line_consumer=None,
497                  timeout=None,
498                  shell=True,
499                  join_stderr=True,
500                  env=None,
501                  except_handler=lambda p, e: None):
502    """Run the command with an extra feature line_consumer.
503
504    This version allow developers to provide a line_consumer which will be
505    fed execution output lines.
506
507    A line_consumer is a callback, which is given a chance to run for each
508    line the execution outputs (either to stdout or stderr). The
509    line_consumer must accept one and exactly one dict argument, the dict
510    argument has these items -
511      'line'   -  The line output by the binary. Notice, this string includes
512                  the trailing '\n'.
513      'output' -  Whether this is a stdout or stderr output, values are either
514                  'stdout' or 'stderr'. When join_stderr is True, this value
515                  will always be 'output'.
516      'pobject' - The object used to control execution, for example, call
517                  pobject.kill().
518
519    Note: As this is written, the stdin for the process executed is
520    not associated with the stdin of the caller of this routine.
521
522    Args:
523      cmd: Command in a single string.
524      cwd: Working directory for execution.
525      line_consumer: A function that will ba called by this function. See above
526        for details.
527      timeout: terminate command after this timeout.
528      shell: Whether to use a shell for execution.
529      join_stderr: Whether join stderr to stdout stream.
530      env: Execution environment.
531      except_handler: Callback for when exception is thrown during command
532        execution. Passed process object and exception.
533
534    Returns:
535      Execution return code.
536
537    Raises:
538      child_exception: if fails to start the command process (missing
539                       permission, no such file, etc)
540    """
541
542    class StreamHandler(object):
543      """Internal utility class."""
544
545      def __init__(self, pobject, fd, name, line_consumer):
546        self._pobject = pobject
547        self._fd = fd
548        self._name = name
549        self._buf = ''
550        self._line_consumer = line_consumer
551
552      def read_and_notify_line(self):
553        t = os.read(fd, 1024)
554        self._buf = self._buf + t
555        self.notify_line()
556
557      def notify_line(self):
558        p = self._buf.find('\n')
559        while p >= 0:
560          self._line_consumer(line=self._buf[:p + 1],
561                              output=self._name,
562                              pobject=self._pobject)
563          if p < len(self._buf) - 1:
564            self._buf = self._buf[p + 1:]
565            p = self._buf.find('\n')
566          else:
567            self._buf = ''
568            p = -1
569            break
570
571      def notify_eos(self):
572        # Notify end of stream. The last line may not end with a '\n'.
573        if self._buf != '':
574          self._line_consumer(line=self._buf,
575                              output=self._name,
576                              pobject=self._pobject)
577          self._buf = ''
578
579    if self.log_level == 'verbose':
580      self.logger.LogCmd(cmd)
581    elif self.logger:
582      self.logger.LogCmdToFileOnly(cmd)
583
584    # We use setsid so that the child will have a different session id
585    # and we can easily kill the process group. This is also important
586    # because the child will be disassociated from the parent terminal.
587    # In this way the child cannot mess the parent's terminal.
588    pobject = None
589    try:
590      pobject = subprocess.Popen(
591          cmd,
592          cwd=cwd,
593          bufsize=1024,
594          env=env,
595          shell=shell,
596          universal_newlines=True,
597          stdout=subprocess.PIPE,
598          stderr=subprocess.STDOUT if join_stderr else subprocess.PIPE,
599          preexec_fn=os.setsid)
600
601      # We provide a default line_consumer
602      if line_consumer is None:
603        line_consumer = lambda **d: None
604      start_time = time.time()
605      poll = select.poll()
606      outfd = pobject.stdout.fileno()
607      poll.register(outfd, select.POLLIN | select.POLLPRI)
608      handlermap = {outfd:
609                    StreamHandler(pobject, outfd, 'stdout', line_consumer)}
610      if not join_stderr:
611        errfd = pobject.stderr.fileno()
612        poll.register(errfd,
613                      select.POLLIN | select.POLLPRI)
614        handlermap[errfd] = StreamHandler(pobject,
615                                          errfd,
616                                          'stderr',
617                                          line_consumer)
618      while len(handlermap):
619        readables = poll.poll(300)
620        for (fd, evt) in readables:
621          handler = handlermap[fd]
622          if evt & (select.POLLPRI | select.POLLIN):
623            handler.read_and_notify_line()
624          elif evt & (select.POLLHUP | select.POLLERR | select.POLLNVAL):
625            handler.notify_eos()
626            poll.unregister(fd)
627            del handlermap[fd]
628
629        if timeout is not None and (time.time() - start_time > timeout):
630          os.killpg(os.getpgid(pobject.pid), signal.SIGTERM)
631
632      return pobject.wait()
633    except BaseException as e:
634      except_handler(pobject, e)
635      raise
636
637
638class MockCommandExecuter(CommandExecuter):
639  """Mock class for class CommandExecuter."""
640
641  def __init__(self, log_level, logger_to_set=None):
642    super(MockCommandExecuter, self).__init__(log_level, logger_to_set)
643
644  def RunCommandGeneric(self,
645                        cmd,
646                        return_output=False,
647                        machine=None,
648                        username=None,
649                        command_terminator=None,
650                        command_timeout=None,
651                        terminated_timeout=10,
652                        print_to_console=True,
653                        except_handler=lambda p, e: None):
654    assert not command_timeout
655    cmd = str(cmd)
656    if machine is None:
657      machine = 'localhost'
658    if username is None:
659      username = 'current'
660    logger.GetLogger().LogCmd('(Mock) ' + cmd, machine, username,
661                              print_to_console)
662    return (0, '', '')
663
664  def RunCommand(self, *args, **kwargs):
665    assert 'return_output' not in kwargs
666    kwargs['return_output'] = False
667    return self.RunCommandGeneric(*args, **kwargs)[0]
668
669  def RunCommandWOutput(self, *args, **kwargs):
670    assert 'return_output' not in kwargs
671    kwargs['return_output'] = True
672    return self.RunCommandGeneric(*args, **kwargs)
673
674
675class CommandTerminator(object):
676  """Object to request termination of a command in execution."""
677
678  def __init__(self):
679    self.terminated = False
680
681  def Terminate(self):
682    self.terminated = True
683
684  def IsTerminated(self):
685    return self.terminated
686