1# Copyright (c) 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
5import logging
6import os
7import pipes
8import threading
9
10from autotest_lib.client.common_lib import error
11
12class _HelperThread(threading.Thread):
13    """Make a thread to run the command in."""
14    def __init__(self, host, cmd):
15        super(_HelperThread, self).__init__()
16        self._host = host
17        self._cmd = cmd
18        self._result = None
19        self.daemon = True
20
21
22    def run(self):
23        logging.info('Helper thread running: %s', self._cmd)
24        # NB: set ignore_status as we're always terminated w/ pkill
25        self._result = self._host.run(self._cmd, ignore_status=True)
26
27
28    @property
29    def result(self):
30        """
31        @returns string result of running our command if the command has
32                finished, and None otherwise.
33
34        """
35        return self._result
36
37
38class Command(object):
39    """
40    Encapsulates a command run on a remote machine.
41
42    Future work is to have this get the PID (by prepending 'echo $$;
43    exec' to the command and parsing the output).
44
45    """
46    def __init__(self, host, cmd, pkill_argument=None):
47        """
48        Run a command on a remote host in the background.
49
50        @param host Host object representing the remote machine.
51        @param cmd String command to run on the remote machine.
52        @param pkill_argument String argument to pkill to kill the remote
53                process.
54
55        """
56        if pkill_argument is None:
57            # Attempt to guess what a suitable pkill argument would look like.
58            pkill_argument = os.path.basename(cmd.split()[0])
59        self._command_name = pipes.quote(pkill_argument)
60        self._host = host
61        self._thread = _HelperThread(self._host, cmd)
62        self._thread.start()
63
64
65    def join(self, signal=None, timeout=5.0):
66        """
67        Kills the remote command and waits until it dies.  Takes an optional
68        signal argument to control which signal to send the process to be
69        killed.
70
71        @param signal Signal string to give to pkill (e.g. SIGNAL_INT).
72        @param timeout float number of seconds to wait for join to finish.
73
74        """
75        if signal is None:
76            signal_arg = ''
77        else:
78            # In theory, it should be hard to pass something evil for signal if
79            # we make sure it's an integer before passing it to pkill.
80            signal_arg = '-' + str(int(signal))
81
82        # Ignore status because the command may have exited already
83        self._host.run("pkill %s %s" % (signal_arg, self._command_name),
84                       ignore_status=True)
85        self._thread.join(timeout)
86        if self._thread.isAlive():
87            raise error.TestFail('Failed to kill remote command: %s' %
88                                 self._command_name)
89
90
91    def __enter__(self):
92        return self
93
94
95    def __exit__(self, exception, value, traceback):
96        self.join()
97        return False
98
99
100    @property
101    def result(self):
102        """
103        @returns string result of running our command if the command has
104                finished, and None otherwise.
105
106        """
107        return self._thread.result
108