• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Lint as: python2, python3
2# Copyright 2016 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""This module provides an object to record the output of command-line program.
7"""
8
9from __future__ import absolute_import
10from __future__ import division
11from __future__ import print_function
12import fcntl
13import logging
14import os
15import pty
16import re
17import subprocess
18import threading
19import time
20
21
22class OutputRecorderError(Exception):
23    """An exception class for output_recorder module."""
24    pass
25
26
27class OutputRecorder(object):
28    """A class used to record the output of command line program.
29
30    A thread is dedicated to performing non-blocking reading of the
31    command outpt in this class. Other possible approaches include
32    1. using gobject.io_add_watch() to register a callback and
33       reading the output when available, or
34    2. using select.select() with a short timeout, and reading
35       the output if available.
36    However, the above two approaches are not very reliable. Hence,
37    this approach using non-blocking reading is adopted.
38
39    To prevent the block buffering of the command output, a pseudo
40    terminal is created through pty.openpty(). This forces the
41    line output.
42
43    This class saves the output in self.contents so that it is
44    easy to perform regular expression search(). The output is
45    also saved in a file.
46
47    """
48
49    DEFAULT_OPEN_MODE = 'a'
50    START_DELAY_SECS = 1        # Delay after starting recording.
51    STOP_DELAY_SECS = 1         # Delay before stopping recording.
52    POLLING_DELAY_SECS = 0.1    # Delay before next polling.
53    TMP_FILE = '/tmp/output_recorder.dat'
54
55    def __init__(self, cmd, open_mode=DEFAULT_OPEN_MODE,
56                 start_delay_secs=START_DELAY_SECS,
57                 stop_delay_secs=STOP_DELAY_SECS, save_file=TMP_FILE):
58        """Construction of output recorder.
59
60        @param cmd: the command of which the output is to record.
61        @param open_mode: the open mode for writing output to save_file.
62                Could be either 'w' or 'a'.
63        @param stop_delay_secs: the delay time before stopping the cmd.
64        @param save_file: the file to save the output.
65
66        """
67        self.cmd = cmd
68        self.open_mode = open_mode
69        self.start_delay_secs = start_delay_secs
70        self.stop_delay_secs = stop_delay_secs
71        self.save_file = save_file
72        self.contents = []
73
74        # Create a thread dedicated to record the output.
75        self._recording_thread = None
76        self._stop_recording_thread_event = threading.Event()
77
78        # Use pseudo terminal to prevent buffering of the program output.
79        self._main, self._node = pty.openpty()
80        self._output = os.fdopen(self._main)
81
82        # Set non-blocking flag.
83        fcntl.fcntl(self._output, fcntl.F_SETFL, os.O_NONBLOCK)
84
85
86    def record(self):
87        """Record the output of the cmd."""
88        logging.info('Recording output of "%s".', self.cmd)
89        try:
90            self._recorder = subprocess.Popen(
91                    self.cmd, stdout=self._node, stderr=self._node)
92        except:
93            raise OutputRecorderError('Failed to run "%s"' % self.cmd)
94
95        with open(self.save_file, self.open_mode) as output_f:
96            output_f.write(os.linesep + '*' * 80 + os.linesep)
97            while True:
98                try:
99                    # Perform non-blocking read.
100                    line = self._output.readline()
101                except:
102                    # Set empty string if nothing to read.
103                    line = ''
104
105                if line:
106                    output_f.write(line)
107                    output_f.flush()
108                    # The output, e.g. the output of btmon, may contain some
109                    # special unicode such that we would like to escape.
110                    # In this way, regular expression search could be conducted
111                    # properly.
112                    self.contents.append(line.encode('unicode-escape'))
113                elif self._stop_recording_thread_event.is_set():
114                    self._stop_recording_thread_event.clear()
115                    break
116                else:
117                    # Sleep a while if nothing to read yet.
118                    time.sleep(self.POLLING_DELAY_SECS)
119
120
121    def start(self):
122        """Start the recording thread."""
123        logging.info('Start recording thread.')
124        self.clear_contents()
125        self._recording_thread = threading.Thread(target=self.record)
126        self._recording_thread.start()
127        time.sleep(self.start_delay_secs)
128
129
130    def stop(self):
131        """Stop the recording thread."""
132        logging.info('Stop recording thread.')
133        time.sleep(self.stop_delay_secs)
134        self._stop_recording_thread_event.set()
135        self._recording_thread.join()
136
137        # Kill the process.
138        self._recorder.terminate()
139        self._recorder.kill()
140
141
142    def clear_contents(self):
143        """Clear the contents."""
144        self.contents = []
145
146
147    def get_contents(self, search_str='', start_str=''):
148        """Get the (filtered) contents.
149
150        @param search_str: only lines with search_str would be kept.
151        @param start_str: all lines before the occurrence of start_str would be
152                          filtered.
153
154        @returns: the (filtered) contents.
155
156        """
157        search_pattern = re.compile(search_str) if search_str else None
158        start_pattern = re.compile(start_str) if start_str else None
159
160        # Just returns the original contents if no filtered conditions are
161        # specified.
162        if not search_pattern and not start_pattern:
163            return self.contents
164
165        contents = []
166        start_flag = not bool(start_pattern)
167        for line in self.contents:
168            if start_flag:
169                if search_pattern.search(line):
170                    contents.append(line.strip())
171            elif start_pattern.search(line):
172                start_flag = True
173                contents.append(line.strip())
174
175        return contents
176
177
178    def find(self, pattern_str, flags=re.I):
179        """Find a pattern string in the contents.
180
181        Note that the pattern_str is considered as an arbitrary literal string
182        that might contain re meta-characters, e.g., '(' or ')'. Hence,
183        re.escape() is applied before using re.compile.
184
185        @param pattern_str: the pattern string to search.
186        @param flags: the flags of the pattern expression behavior.
187
188        @returns: True if found. False otherwise.
189
190        """
191        pattern = re.compile(re.escape(pattern_str), flags)
192        for line in self.contents:
193            result = pattern.search(line)
194            if result:
195                return True
196        return False
197
198
199if __name__ == '__main__':
200    # A demo using btmon tool to monitor bluetoohd activity.
201    cmd = 'btmon'
202    recorder = OutputRecorder(cmd)
203
204    if True:
205        recorder.start()
206        # Perform some bluetooth activities here in another terminal.
207        time.sleep(recorder.stop_delay_secs)
208        recorder.stop()
209
210    for line in recorder.get_contents():
211        print(line)
212