blob: 9222d3302a157aa89a95a05533b589491cad8df7 [file] [log] [blame]
lmrbbc9dd52009-07-22 20:33:47 +00001#!/usr/bin/python
lmrbbc9dd52009-07-22 20:33:47 +00002"""
3A class and functions used for running and controlling child processes.
4
5@copyright: 2008-2009 Red Hat Inc.
6"""
7
lmr78fabe52009-10-05 18:52:51 +00008import os, sys, pty, select, termios, fcntl
9
10
11# The following helper functions are shared by the server and the client.
12
13def _lock(filename):
14 if not os.path.exists(filename):
15 open(filename, "w").close()
16 fd = os.open(filename, os.O_RDWR)
17 fcntl.lockf(fd, fcntl.LOCK_EX)
18 return fd
19
20
21def _unlock(fd):
22 fcntl.lockf(fd, fcntl.LOCK_UN)
23 os.close(fd)
24
25
26def _locked(filename):
27 try:
28 fd = os.open(filename, os.O_RDWR)
29 except:
30 return False
31 try:
32 fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
33 except:
34 os.close(fd)
35 return True
36 fcntl.lockf(fd, fcntl.LOCK_UN)
37 os.close(fd)
38 return False
39
40
41def _wait(filename):
42 fd = _lock(filename)
43 _unlock(fd)
44
45
46def _get_filenames(base_dir, id):
47 return [os.path.join(base_dir, s + id) for s in
48 "shell-pid-", "status-", "output-", "inpipe-",
49 "lock-server-running-", "lock-client-starting-"]
50
51
52def _get_reader_filename(base_dir, id, reader):
53 return os.path.join(base_dir, "outpipe-%s-%s" % (reader, id))
54
55
56# The following is the server part of the module.
57
58if __name__ == "__main__":
59 id = sys.stdin.readline().strip()
60 echo = sys.stdin.readline().strip() == "True"
61 readers = sys.stdin.readline().strip().split(",")
62 command = sys.stdin.readline().strip() + " && echo %s > /dev/null" % id
63
64 # Define filenames to be used for communication
65 base_dir = "/tmp/kvm_spawn"
66 (shell_pid_filename,
67 status_filename,
68 output_filename,
69 inpipe_filename,
70 lock_server_running_filename,
71 lock_client_starting_filename) = _get_filenames(base_dir, id)
72
73 # Populate the reader filenames list
74 reader_filenames = [_get_reader_filename(base_dir, id, reader)
75 for reader in readers]
76
77 # Set $TERM = dumb
78 os.putenv("TERM", "dumb")
79
80 (shell_pid, shell_fd) = pty.fork()
81 if shell_pid == 0:
82 # Child process: run the command in a subshell
83 os.execv("/bin/sh", ["/bin/sh", "-c", command])
84 else:
85 # Parent process
86 lock_server_running = _lock(lock_server_running_filename)
87
88 # Set terminal echo on/off and disable pre- and post-processing
89 attr = termios.tcgetattr(shell_fd)
90 attr[0] &= ~termios.INLCR
91 attr[0] &= ~termios.ICRNL
92 attr[0] &= ~termios.IGNCR
93 attr[1] &= ~termios.OPOST
94 if echo:
95 attr[3] |= termios.ECHO
96 else:
97 attr[3] &= ~termios.ECHO
98 termios.tcsetattr(shell_fd, termios.TCSANOW, attr)
99
100 # Open output file
101 output_file = open(output_filename, "w")
102 # Open input pipe
103 os.mkfifo(inpipe_filename)
104 inpipe_fd = os.open(inpipe_filename, os.O_RDWR)
105 # Open output pipes (readers)
106 reader_fds = []
107 for filename in reader_filenames:
108 os.mkfifo(filename)
109 reader_fds.append(os.open(filename, os.O_RDWR))
110
111 # Write shell PID to file
112 file = open(shell_pid_filename, "w")
113 file.write(str(shell_pid))
114 file.close()
115
116 # Print something to stdout so the client can start working
117 print "Server %s ready" % id
118 sys.stdout.flush()
119
120 # Initialize buffers
121 buffers = ["" for reader in readers]
122
123 # Read from child and write to files/pipes
124 while True:
125 check_termination = False
126 # Make a list of reader pipes whose buffers are not empty
127 fds = [fd for (i, fd) in enumerate(reader_fds) if buffers[i]]
128 # Wait until there's something to do
129 r, w, x = select.select([shell_fd, inpipe_fd], fds, [], 0.5)
130 # If a reader pipe is ready for writing --
131 for (i, fd) in enumerate(reader_fds):
132 if fd in w:
133 bytes_written = os.write(fd, buffers[i])
134 buffers[i] = buffers[i][bytes_written:]
135 # If there's data to read from the child process --
136 if shell_fd in r:
137 try:
138 data = os.read(shell_fd, 16384)
139 except OSError:
140 data = ""
141 if not data:
142 check_termination = True
143 # Remove carriage returns from the data -- they often cause
144 # trouble and are normally not needed
145 data = data.replace("\r", "")
146 output_file.write(data)
147 output_file.flush()
148 for i in range(len(readers)):
149 buffers[i] += data
150 # If os.read() raised an exception or there was nothing to read --
151 if check_termination or shell_fd not in r:
152 pid, status = os.waitpid(shell_pid, os.WNOHANG)
153 if pid:
154 status = os.WEXITSTATUS(status)
155 break
156 # If there's data to read from the client --
157 if inpipe_fd in r:
158 data = os.read(inpipe_fd, 1024)
159 os.write(shell_fd, data)
160
161 # Write the exit status to a file
162 file = open(status_filename, "w")
163 file.write(str(status))
164 file.close()
165
166 # Wait for the client to finish initializing
167 _wait(lock_client_starting_filename)
168
169 # Delete FIFOs
170 for filename in reader_filenames + [inpipe_filename]:
171 try:
172 os.unlink(filename)
173 except OSError:
174 pass
175
176 # Close all files and pipes
177 output_file.close()
178 os.close(inpipe_fd)
179 for fd in reader_fds:
180 os.close(fd)
181
182 _unlock(lock_server_running)
183 exit(0)
184
185
186# The following is the client part of the module.
187
188import subprocess, time, signal, re, threading, logging
lmrb635b862009-09-10 14:53:21 +0000189import common, kvm_utils
190
lmrbbc9dd52009-07-22 20:33:47 +0000191
192def run_bg(command, termination_func=None, output_func=None, output_prefix="",
193 timeout=1.0):
194 """
195 Run command as a subprocess. Call output_func with each line of output
196 from the subprocess (prefixed by output_prefix). Call termination_func
197 when the subprocess terminates. Return when timeout expires or when the
198 subprocess exits -- whichever occurs first.
199
200 @brief: Run a subprocess in the background and collect its output and
201 exit status.
202
203 @param command: The shell command to execute
204 @param termination_func: A function to call when the process terminates
205 (should take an integer exit status parameter)
206 @param output_func: A function to call with each line of output from
207 the subprocess (should take a string parameter)
208 @param output_prefix: A string to pre-pend to each line of the output,
209 before passing it to stdout_func
210 @param timeout: Time duration (in seconds) to wait for the subprocess to
211 terminate before returning
212
213 @return: A kvm_tail object.
214 """
215 process = kvm_tail(command=command,
216 termination_func=termination_func,
217 output_func=output_func,
218 output_prefix=output_prefix)
219
220 end_time = time.time() + timeout
221 while time.time() < end_time and process.is_alive():
222 time.sleep(0.1)
223
224 return process
225
226
227def run_fg(command, output_func=None, output_prefix="", timeout=1.0):
228 """
229 Run command as a subprocess. Call output_func with each line of output
230 from the subprocess (prefixed by prefix). Return when timeout expires or
231 when the subprocess exits -- whichever occurs first. If timeout expires
232 and the subprocess is still running, kill it before returning.
233
234 @brief: Run a subprocess in the foreground and collect its output and
235 exit status.
236
237 @param command: The shell command to execute
238 @param output_func: A function to call with each line of output from
239 the subprocess (should take a string parameter)
240 @param output_prefix: A string to pre-pend to each line of the output,
241 before passing it to stdout_func
242 @param timeout: Time duration (in seconds) to wait for the subprocess to
243 terminate before killing it and returning
244
245 @return: A 2-tuple containing the exit status of the process and its
246 STDOUT/STDERR output. If timeout expires before the process
247 terminates, the returned status is None.
248 """
249 process = run_bg(command, None, output_func, output_prefix, timeout)
250 output = process.get_output()
251 if process.is_alive():
252 status = None
253 else:
254 status = process.get_status()
255 process.close()
256 return (status, output)
257
258
lmrbbc9dd52009-07-22 20:33:47 +0000259class kvm_spawn:
260 """
261 This class is used for spawning and controlling a child process.
262
263 A new instance of this class can either run a new server (a small Python
264 program that reads output from the child process and reports it to the
265 client and to a text file) or attach to an already running server.
266 When a server is started it runs the child process.
267 The server writes output from the child's STDOUT and STDERR to a text file.
268 The text file can be accessed at any time using get_output().
269 In addition, the server opens as many pipes as requested by the client and
270 writes the output to them.
271 The pipes are requested and accessed by classes derived from kvm_spawn.
272 These pipes are referred to as "readers".
273 The server also receives input from the client and sends it to the child
274 process.
275 An instance of this class can be pickled. Every derived class is
276 responsible for restoring its own state by properly defining
277 __getinitargs__().
278
279 The first named pipe is used by _tail(), a function that runs in the
280 background and reports new output from the child as it is produced.
281 The second named pipe is used by a set of functions that read and parse
282 output as requested by the user in an interactive manner, similar to
283 pexpect.
284 When unpickled it automatically
285 resumes _tail() if needed.
286 """
287
288 def __init__(self, command=None, id=None, echo=False, linesep="\n"):
289 """
290 Initialize the class and run command as a child process.
291
292 @param command: Command to run, or None if accessing an already running
293 server.
294 @param id: ID of an already running server, if accessing a running
295 server, or None if starting a new one.
296 @param echo: Boolean indicating whether echo should be initially
297 enabled for the pseudo terminal running the subprocess. This
298 parameter has an effect only when starting a new server.
299 @param linesep: Line separator to be appended to strings sent to the
300 child process by sendline().
301 """
302 self.id = id or kvm_utils.generate_random_string(8)
303
304 # Define filenames for communication with server
305 base_dir = "/tmp/kvm_spawn"
306 try:
307 os.makedirs(base_dir)
308 except:
309 pass
310 (self.shell_pid_filename,
311 self.status_filename,
312 self.output_filename,
313 self.inpipe_filename,
314 self.lock_server_running_filename,
315 self.lock_client_starting_filename) = _get_filenames(base_dir,
316 self.id)
317
318 # Remember some attributes
319 self.echo = echo
320 self.linesep = linesep
321
322 # Make sure the 'readers' and 'close_hooks' attributes exist
323 if not hasattr(self, "readers"):
324 self.readers = []
325 if not hasattr(self, "close_hooks"):
326 self.close_hooks = []
327
328 # Define the reader filenames
329 self.reader_filenames = dict(
330 (reader, _get_reader_filename(base_dir, self.id, reader))
331 for reader in self.readers)
332
333 # Let the server know a client intends to open some pipes;
334 # if the executed command terminates quickly, the server will wait for
335 # the client to release the lock before exiting
336 lock_client_starting = _lock(self.lock_client_starting_filename)
337
338 # Start the server (which runs the command)
339 if command:
340 sub = subprocess.Popen("python %s" % __file__,
341 shell=True,
342 stdin=subprocess.PIPE,
343 stdout=subprocess.PIPE,
344 stderr=subprocess.STDOUT)
345 # Send parameters to the server
346 sub.stdin.write("%s\n" % self.id)
347 sub.stdin.write("%s\n" % echo)
348 sub.stdin.write("%s\n" % ",".join(self.readers))
349 sub.stdin.write("%s\n" % command)
350 # Wait for the server to complete its initialization
lmrb8f53d62009-07-27 13:29:17 +0000351 while not "Server %s ready" % self.id in sub.stdout.readline():
352 pass
lmrbbc9dd52009-07-22 20:33:47 +0000353
354 # Open the reading pipes
355 self.reader_fds = {}
356 try:
357 assert(_locked(self.lock_server_running_filename))
358 for reader, filename in self.reader_filenames.items():
359 self.reader_fds[reader] = os.open(filename, os.O_RDONLY)
360 except:
361 pass
362
363 # Allow the server to continue
364 _unlock(lock_client_starting)
365
366
367 # The following two functions are defined to make sure the state is set
368 # exclusively by the constructor call as specified in __getinitargs__().
369
370 def __getstate__(self):
371 pass
372
373
374 def __setstate__(self, state):
375 pass
376
377
378 def __getinitargs__(self):
379 # Save some information when pickling -- will be passed to the
380 # constructor upon unpickling
381 return (None, self.id, self.echo, self.linesep)
382
383
384 def _add_reader(self, reader):
385 """
386 Add a reader whose file descriptor can be obtained with _get_fd().
387 Should be called before __init__(). Intended for use by derived
388 classes.
389
390 @param reader: The name of the reader.
391 """
392 if not hasattr(self, "readers"):
393 self.readers = []
394 self.readers.append(reader)
395
396
397 def _add_close_hook(self, hook):
398 """
399 Add a close hook function to be called when close() is called.
400 The function will be called after the process terminates but before
401 final cleanup. Intended for use by derived classes.
402
403 @param hook: The hook function.
404 """
405 if not hasattr(self, "close_hooks"):
406 self.close_hooks = []
407 self.close_hooks.append(hook)
408
409
410 def _get_fd(self, reader):
411 """
412 Return an open file descriptor corresponding to the specified reader
413 pipe. If no such reader exists, or the pipe could not be opened,
414 return None. Intended for use by derived classes.
415
416 @param reader: The name of the reader.
417 """
418 return self.reader_fds.get(reader)
419
420
421 def get_id(self):
422 """
423 Return the instance's id attribute, which may be used to access the
424 process in the future.
425 """
426 return self.id
427
428
lmrfb151b52009-09-09 22:19:11 +0000429 def get_pid(self):
lmrbbc9dd52009-07-22 20:33:47 +0000430 """
lmrfb151b52009-09-09 22:19:11 +0000431 Return the PID of the process.
432
433 Note: this may be the PID of the shell process running the user given
434 command.
lmrbbc9dd52009-07-22 20:33:47 +0000435 """
436 try:
437 file = open(self.shell_pid_filename, "r")
438 pid = int(file.read())
439 file.close()
440 return pid
441 except:
442 return None
443
444
lmrbbc9dd52009-07-22 20:33:47 +0000445 def get_status(self):
446 """
447 Wait for the process to exit and return its exit status, or None
448 if the exit status is not available.
449 """
450 _wait(self.lock_server_running_filename)
451 try:
452 file = open(self.status_filename, "r")
453 status = int(file.read())
454 file.close()
455 return status
456 except:
457 return None
458
459
460 def get_output(self):
461 """
462 Return the STDOUT and STDERR output of the process so far.
463 """
464 try:
465 file = open(self.output_filename, "r")
466 output = file.read()
467 file.close()
468 return output
469 except:
470 return ""
471
472
473 def is_alive(self):
474 """
475 Return True if the process is running.
476 """
lmr5df99f32009-08-13 04:46:16 +0000477 return _locked(self.lock_server_running_filename)
lmrbbc9dd52009-07-22 20:33:47 +0000478
479
lmrea1b64d2009-09-09 22:14:09 +0000480 def close(self, sig=signal.SIGKILL):
lmrbbc9dd52009-07-22 20:33:47 +0000481 """
482 Kill the child process if it's alive and remove temporary files.
483
484 @param sig: The signal to send the process when attempting to kill it.
485 """
486 # Kill it if it's alive
487 if self.is_alive():
lmrfb151b52009-09-09 22:19:11 +0000488 kvm_utils.kill_process_tree(self.get_pid(), sig)
lmrbbc9dd52009-07-22 20:33:47 +0000489 # Wait for the server to exit
490 _wait(self.lock_server_running_filename)
491 # Call all cleanup routines
492 for hook in self.close_hooks:
lmr16063962009-10-14 10:27:59 +0000493 hook(self)
lmrbbc9dd52009-07-22 20:33:47 +0000494 # Close reader file descriptors
495 for fd in self.reader_fds.values():
496 try:
497 os.close(fd)
498 except:
499 pass
lmr04d5b012009-11-10 16:28:22 +0000500 self.reader_fds = {}
lmrbbc9dd52009-07-22 20:33:47 +0000501 # Remove all used files
502 for filename in (_get_filenames("/tmp/kvm_spawn", self.id) +
503 self.reader_filenames.values()):
504 try:
505 os.unlink(filename)
506 except OSError:
507 pass
508
509
510 def set_linesep(self, linesep):
511 """
512 Sets the line separator string (usually "\\n").
513
514 @param linesep: Line separator string.
515 """
516 self.linesep = linesep
517
518
519 def send(self, str=""):
520 """
521 Send a string to the child process.
522
523 @param str: String to send to the child process.
524 """
525 try:
526 fd = os.open(self.inpipe_filename, os.O_RDWR)
527 os.write(fd, str)
528 os.close(fd)
529 except:
530 pass
531
532
533 def sendline(self, str=""):
534 """
535 Send a string followed by a line separator to the child process.
536
537 @param str: String to send to the child process.
538 """
539 self.send(str + self.linesep)
540
541
542class kvm_tail(kvm_spawn):
543 """
544 This class runs a child process in the background and sends its output in
545 real time, line-by-line, to a callback function.
546
547 See kvm_spawn's docstring.
548
549 This class uses a single pipe reader to read data in real time from the
550 child process and report it to a given callback function.
551 When the child process exits, its exit status is reported to an additional
552 callback function.
553
554 When this class is unpickled, it automatically resumes reporting output.
555 """
556
557 def __init__(self, command=None, id=None, echo=False, linesep="\n",
lmr4fe15ba2009-08-13 04:11:26 +0000558 termination_func=None, termination_params=(),
559 output_func=None, output_params=(),
560 output_prefix=""):
lmrbbc9dd52009-07-22 20:33:47 +0000561 """
562 Initialize the class and run command as a child process.
563
564 @param command: Command to run, or None if accessing an already running
565 server.
566 @param id: ID of an already running server, if accessing a running
567 server, or None if starting a new one.
568 @param echo: Boolean indicating whether echo should be initially
569 enabled for the pseudo terminal running the subprocess. This
570 parameter has an effect only when starting a new server.
571 @param linesep: Line separator to be appended to strings sent to the
572 child process by sendline().
573 @param termination_func: Function to call when the process exits. The
574 function must accept a single exit status parameter.
lmr4fe15ba2009-08-13 04:11:26 +0000575 @param termination_params: Parameters to send to termination_func
576 before the exit status.
lmrbbc9dd52009-07-22 20:33:47 +0000577 @param output_func: Function to call whenever a line of output is
578 available from the STDOUT or STDERR streams of the process.
579 The function must accept a single string parameter. The string
580 does not include the final newline.
lmr4fe15ba2009-08-13 04:11:26 +0000581 @param output_params: Parameters to send to output_func before the
582 output line.
lmrbbc9dd52009-07-22 20:33:47 +0000583 @param output_prefix: String to prepend to lines sent to output_func.
584 """
585 # Add a reader and a close hook
586 self._add_reader("tail")
lmr16063962009-10-14 10:27:59 +0000587 self._add_close_hook(kvm_tail._join_thread)
lmrbbc9dd52009-07-22 20:33:47 +0000588
589 # Init the superclass
590 kvm_spawn.__init__(self, command, id, echo, linesep)
591
592 # Remember some attributes
593 self.termination_func = termination_func
lmr4fe15ba2009-08-13 04:11:26 +0000594 self.termination_params = termination_params
lmrbbc9dd52009-07-22 20:33:47 +0000595 self.output_func = output_func
lmr4fe15ba2009-08-13 04:11:26 +0000596 self.output_params = output_params
lmrbbc9dd52009-07-22 20:33:47 +0000597 self.output_prefix = output_prefix
598
599 # Start the thread in the background
lmrc07f3812009-10-14 10:27:04 +0000600 self.tail_thread = None
lmra4197002009-08-13 05:00:51 +0000601 self.__thread_kill_requested = False
lmrc07f3812009-10-14 10:27:04 +0000602 if termination_func or output_func:
603 self._start_thread()
lmrbbc9dd52009-07-22 20:33:47 +0000604
605
606 def __getinitargs__(self):
607 return kvm_spawn.__getinitargs__(self) + (self.termination_func,
lmr4fe15ba2009-08-13 04:11:26 +0000608 self.termination_params,
lmrbbc9dd52009-07-22 20:33:47 +0000609 self.output_func,
lmr4fe15ba2009-08-13 04:11:26 +0000610 self.output_params,
lmrbbc9dd52009-07-22 20:33:47 +0000611 self.output_prefix)
612
613
614 def set_termination_func(self, termination_func):
615 """
616 Set the termination_func attribute. See __init__() for details.
617
618 @param termination_func: Function to call when the process terminates.
619 Must take a single parameter -- the exit status.
620 """
621 self.termination_func = termination_func
lmrc07f3812009-10-14 10:27:04 +0000622 if termination_func and not self.tail_thread:
623 self._start_thread()
lmrbbc9dd52009-07-22 20:33:47 +0000624
625
lmr4fe15ba2009-08-13 04:11:26 +0000626 def set_termination_params(self, termination_params):
627 """
628 Set the termination_params attribute. See __init__() for details.
629
630 @param termination_params: Parameters to send to termination_func
631 before the exit status.
632 """
633 self.termination_params = termination_params
634
635
lmrbbc9dd52009-07-22 20:33:47 +0000636 def set_output_func(self, output_func):
637 """
638 Set the output_func attribute. See __init__() for details.
639
640 @param output_func: Function to call for each line of STDOUT/STDERR
641 output from the process. Must take a single string parameter.
642 """
643 self.output_func = output_func
lmrc07f3812009-10-14 10:27:04 +0000644 if output_func and not self.tail_thread:
645 self._start_thread()
lmrbbc9dd52009-07-22 20:33:47 +0000646
647
lmr4fe15ba2009-08-13 04:11:26 +0000648 def set_output_params(self, output_params):
649 """
650 Set the output_params attribute. See __init__() for details.
651
652 @param output_params: Parameters to send to output_func before the
653 output line.
654 """
655 self.output_params = output_params
656
657
lmrbbc9dd52009-07-22 20:33:47 +0000658 def set_output_prefix(self, output_prefix):
659 """
660 Set the output_prefix attribute. See __init__() for details.
661
662 @param output_prefix: String to pre-pend to each line sent to
663 output_func (see set_output_callback()).
664 """
665 self.output_prefix = output_prefix
666
667
lmra4197002009-08-13 05:00:51 +0000668 def kill_tail_thread(self):
669 """
670 Stop the tailing thread which calls output_func() and
671 termination_func().
672 """
673 self.__thread_kill_requested = True
674 self._join_thread()
675
676
lmrbbc9dd52009-07-22 20:33:47 +0000677 def _tail(self):
678 def print_line(text):
679 # Pre-pend prefix and remove trailing whitespace
680 text = self.output_prefix + text.rstrip()
681 # Sanitize text
682 text = text.decode("utf-8", "replace")
683 # Pass it to output_func
684 try:
lmr4fe15ba2009-08-13 04:11:26 +0000685 params = self.output_params + (text,)
686 self.output_func(*params)
lmrbbc9dd52009-07-22 20:33:47 +0000687 except TypeError:
688 pass
689
690 fd = self._get_fd("tail")
691 buffer = ""
692 while True:
lmra4197002009-08-13 05:00:51 +0000693 if self.__thread_kill_requested:
694 return
lmrbbc9dd52009-07-22 20:33:47 +0000695 try:
696 # See if there's any data to read from the pipe
697 r, w, x = select.select([fd], [], [], 0.05)
698 except:
699 break
700 if fd in r:
701 # Some data is available; read it
702 new_data = os.read(fd, 1024)
703 if not new_data:
704 break
705 buffer += new_data
706 # Send the output to output_func line by line
707 # (except for the last line)
708 if self.output_func:
709 lines = buffer.split("\n")
710 for line in lines[:-1]:
711 print_line(line)
712 # Leave only the last line
713 last_newline_index = buffer.rfind("\n")
714 buffer = buffer[last_newline_index+1:]
715 else:
716 # No output is available right now; flush the buffer
717 if buffer:
718 print_line(buffer)
719 buffer = ""
720 # The process terminated; print any remaining output
721 if buffer:
722 print_line(buffer)
723 # Get the exit status, print it and send it to termination_func
724 status = self.get_status()
725 if status is None:
726 return
727 print_line("(Process terminated with status %s)" % status)
728 try:
lmr4fe15ba2009-08-13 04:11:26 +0000729 params = self.termination_params + (status,)
730 self.termination_func(*params)
lmrbbc9dd52009-07-22 20:33:47 +0000731 except TypeError:
732 pass
733
734
lmrc07f3812009-10-14 10:27:04 +0000735 def _start_thread(self):
736 self.tail_thread = threading.Thread(None, self._tail)
737 self.tail_thread.start()
738
739
lmrbbc9dd52009-07-22 20:33:47 +0000740 def _join_thread(self):
741 # Wait for the tail thread to exit
742 if self.tail_thread:
743 self.tail_thread.join()
744
745
746class kvm_expect(kvm_tail):
747 """
748 This class runs a child process in the background and provides expect-like
749 services.
750
751 It also provides all of kvm_tail's functionality.
752 """
753
754 def __init__(self, command=None, id=None, echo=False, linesep="\n",
lmr4fe15ba2009-08-13 04:11:26 +0000755 termination_func=None, termination_params=(),
756 output_func=None, output_params=(),
757 output_prefix=""):
lmrbbc9dd52009-07-22 20:33:47 +0000758 """
759 Initialize the class and run command as a child process.
760
761 @param command: Command to run, or None if accessing an already running
762 server.
763 @param id: ID of an already running server, if accessing a running
764 server, or None if starting a new one.
765 @param echo: Boolean indicating whether echo should be initially
766 enabled for the pseudo terminal running the subprocess. This
767 parameter has an effect only when starting a new server.
768 @param linesep: Line separator to be appended to strings sent to the
769 child process by sendline().
770 @param termination_func: Function to call when the process exits. The
771 function must accept a single exit status parameter.
lmr4fe15ba2009-08-13 04:11:26 +0000772 @param termination_params: Parameters to send to termination_func
773 before the exit status.
lmrbbc9dd52009-07-22 20:33:47 +0000774 @param output_func: Function to call whenever a line of output is
775 available from the STDOUT or STDERR streams of the process.
776 The function must accept a single string parameter. The string
777 does not include the final newline.
lmr4fe15ba2009-08-13 04:11:26 +0000778 @param output_params: Parameters to send to output_func before the
779 output line.
lmrbbc9dd52009-07-22 20:33:47 +0000780 @param output_prefix: String to prepend to lines sent to output_func.
781 """
782 # Add a reader
783 self._add_reader("expect")
784
785 # Init the superclass
786 kvm_tail.__init__(self, command, id, echo, linesep,
lmr4fe15ba2009-08-13 04:11:26 +0000787 termination_func, termination_params,
788 output_func, output_params, output_prefix)
lmrbbc9dd52009-07-22 20:33:47 +0000789
790
791 def __getinitargs__(self):
792 return kvm_tail.__getinitargs__(self)
793
794
795 def read_nonblocking(self, timeout=None):
796 """
797 Read from child until there is nothing to read for timeout seconds.
798
799 @param timeout: Time (seconds) to wait before we give up reading from
800 the child process, or None to use the default value.
801 """
802 if timeout is None:
803 timeout = 0.1
804 fd = self._get_fd("expect")
805 data = ""
806 while True:
807 try:
808 r, w, x = select.select([fd], [], [], timeout)
809 except:
810 return data
811 if fd in r:
812 new_data = os.read(fd, 1024)
813 if not new_data:
814 return data
815 data += new_data
816 else:
817 return data
818
819
820 def match_patterns(self, str, patterns):
821 """
822 Match str against a list of patterns.
823
824 Return the index of the first pattern that matches a substring of str.
825 None and empty strings in patterns are ignored.
826 If no match is found, return None.
827
828 @param patterns: List of strings (regular expression patterns).
829 """
830 for i in range(len(patterns)):
831 if not patterns[i]:
832 continue
833 if re.search(patterns[i], str):
834 return i
835
836
837 def read_until_output_matches(self, patterns, filter=lambda x: x,
838 timeout=30.0, internal_timeout=None,
839 print_func=None):
840 """
841 Read using read_nonblocking until a match is found using match_patterns,
842 or until timeout expires. Before attempting to search for a match, the
843 data is filtered using the filter function provided.
844
845 @brief: Read from child using read_nonblocking until a pattern
846 matches.
847 @param patterns: List of strings (regular expression patterns)
848 @param filter: Function to apply to the data read from the child before
849 attempting to match it against the patterns (should take and
850 return a string)
851 @param timeout: The duration (in seconds) to wait until a match is
852 found
853 @param internal_timeout: The timeout to pass to read_nonblocking
854 @param print_func: A function to be used to print the data being read
855 (should take a string parameter)
856 @return: Tuple containing the match index (or None if no match was
857 found) and the data read so far.
858 """
859 match = None
860 data = ""
861
lmrf48c5cc2009-10-05 19:22:49 +0000862 fd = self._get_fd("expect")
lmrbbc9dd52009-07-22 20:33:47 +0000863 end_time = time.time() + timeout
lmrf48c5cc2009-10-05 19:22:49 +0000864 while True:
lmr04d5b012009-11-10 16:28:22 +0000865 try:
866 r, w, x = select.select([fd], [], [],
867 max(0, end_time - time.time()))
868 except (select.error, TypeError):
869 break
870 if fd not in r:
871 break
lmrbbc9dd52009-07-22 20:33:47 +0000872 # Read data from child
873 newdata = self.read_nonblocking(internal_timeout)
874 # Print it if necessary
875 if print_func and newdata:
876 str = newdata
877 if str.endswith("\n"):
878 str = str[:-1]
879 for line in str.split("\n"):
880 print_func(line.decode("utf-8", "replace"))
881 data += newdata
882
883 done = False
884 # Look for patterns
885 match = self.match_patterns(filter(data), patterns)
886 if match is not None:
887 done = True
888 # Check if child has died
889 if not self.is_alive():
lmrf48c5cc2009-10-05 19:22:49 +0000890 logging.debug("Process terminated with status %s" %
891 self.get_status())
lmrbbc9dd52009-07-22 20:33:47 +0000892 done = True
893 # Are we done?
894 if done: break
895
896 # Print some debugging info
897 if match is None and (self.is_alive() or self.get_status() != 0):
898 logging.debug("Timeout elapsed or process terminated. Output:" +
899 kvm_utils.format_str_for_message(data.strip()))
900
901 return (match, data)
902
903
904 def read_until_last_word_matches(self, patterns, timeout=30.0,
905 internal_timeout=None, print_func=None):
906 """
907 Read using read_nonblocking until the last word of the output matches
908 one of the patterns (using match_patterns), or until timeout expires.
909
910 @param patterns: A list of strings (regular expression patterns)
911 @param timeout: The duration (in seconds) to wait until a match is
912 found
913 @param internal_timeout: The timeout to pass to read_nonblocking
914 @param print_func: A function to be used to print the data being read
915 (should take a string parameter)
916 @return: A tuple containing the match index (or None if no match was
917 found) and the data read so far.
918 """
919 def get_last_word(str):
920 if str:
921 return str.split()[-1]
922 else:
923 return ""
924
925 return self.read_until_output_matches(patterns, get_last_word,
926 timeout, internal_timeout,
927 print_func)
928
929
930 def read_until_last_line_matches(self, patterns, timeout=30.0,
931 internal_timeout=None, print_func=None):
932 """
933 Read using read_nonblocking until the last non-empty line of the output
934 matches one of the patterns (using match_patterns), or until timeout
935 expires. Return a tuple containing the match index (or None if no match
936 was found) and the data read so far.
937
938 @brief: Read using read_nonblocking until the last non-empty line
939 matches a pattern.
940
941 @param patterns: A list of strings (regular expression patterns)
942 @param timeout: The duration (in seconds) to wait until a match is
943 found
944 @param internal_timeout: The timeout to pass to read_nonblocking
945 @param print_func: A function to be used to print the data being read
946 (should take a string parameter)
947 """
948 def get_last_nonempty_line(str):
949 nonempty_lines = [l for l in str.splitlines() if l.strip()]
950 if nonempty_lines:
951 return nonempty_lines[-1]
952 else:
953 return ""
954
955 return self.read_until_output_matches(patterns, get_last_nonempty_line,
956 timeout, internal_timeout,
957 print_func)
958
959
960class kvm_shell_session(kvm_expect):
961 """
962 This class runs a child process in the background. It it suited for
963 processes that provide an interactive shell, such as SSH and Telnet.
964
965 It provides all services of kvm_expect and kvm_tail. In addition, it
966 provides command running services, and a utility function to test the
967 process for responsiveness.
968 """
969
970 def __init__(self, command=None, id=None, echo=False, linesep="\n",
lmr4fe15ba2009-08-13 04:11:26 +0000971 termination_func=None, termination_params=(),
972 output_func=None, output_params=(),
973 output_prefix="",
lmrbbc9dd52009-07-22 20:33:47 +0000974 prompt=r"[\#\$]\s*$", status_test_command="echo $?"):
975 """
976 Initialize the class and run command as a child process.
977
978 @param command: Command to run, or None if accessing an already running
979 server.
980 @param id: ID of an already running server, if accessing a running
981 server, or None if starting a new one.
982 @param echo: Boolean indicating whether echo should be initially
983 enabled for the pseudo terminal running the subprocess. This
984 parameter has an effect only when starting a new server.
985 @param linesep: Line separator to be appended to strings sent to the
986 child process by sendline().
987 @param termination_func: Function to call when the process exits. The
988 function must accept a single exit status parameter.
lmr4fe15ba2009-08-13 04:11:26 +0000989 @param termination_params: Parameters to send to termination_func
990 before the exit status.
lmrbbc9dd52009-07-22 20:33:47 +0000991 @param output_func: Function to call whenever a line of output is
992 available from the STDOUT or STDERR streams of the process.
993 The function must accept a single string parameter. The string
994 does not include the final newline.
lmr4fe15ba2009-08-13 04:11:26 +0000995 @param output_params: Parameters to send to output_func before the
996 output line.
lmrbbc9dd52009-07-22 20:33:47 +0000997 @param output_prefix: String to prepend to lines sent to output_func.
998 @param prompt: Regular expression describing the shell's prompt line.
999 @param status_test_command: Command to be used for getting the last
1000 exit status of commands run inside the shell (used by
1001 get_command_status_output() and friends).
1002 """
1003 # Init the superclass
1004 kvm_expect.__init__(self, command, id, echo, linesep,
lmr4fe15ba2009-08-13 04:11:26 +00001005 termination_func, termination_params,
1006 output_func, output_params, output_prefix)
lmrbbc9dd52009-07-22 20:33:47 +00001007
1008 # Remember some attributes
1009 self.prompt = prompt
1010 self.status_test_command = status_test_command
1011
1012
1013 def __getinitargs__(self):
1014 return kvm_expect.__getinitargs__(self) + (self.prompt,
1015 self.status_test_command)
1016
1017
lmr97fd76c2009-10-14 10:29:27 +00001018 def __del__(self):
1019 self.close()
1020
1021
lmrbbc9dd52009-07-22 20:33:47 +00001022 def set_prompt(self, prompt):
1023 """
1024 Set the prompt attribute for later use by read_up_to_prompt.
1025
1026 @param: String that describes the prompt contents.
1027 """
1028 self.prompt = prompt
1029
1030
1031 def set_status_test_command(self, status_test_command):
1032 """
1033 Set the command to be sent in order to get the last exit status.
1034
1035 @param status_test_command: Command that will be sent to get the last
1036 exit status.
1037 """
1038 self.status_test_command = status_test_command
1039
1040
1041 def is_responsive(self, timeout=5.0):
1042 """
1043 Return True if the process responds to STDIN/terminal input.
1044
1045 Send a newline to the child process (e.g. SSH or Telnet) and read some
1046 output using read_nonblocking().
1047 If all is OK, some output should be available (e.g. the shell prompt).
1048 In that case return True. Otherwise return False.
1049
1050 @param timeout: Time duration to wait before the process is considered
1051 unresponsive.
1052 """
1053 # Read all output that's waiting to be read, to make sure the output
1054 # we read next is in response to the newline sent
lmr59b98db2009-10-05 19:11:21 +00001055 self.read_nonblocking(timeout=0)
lmrbbc9dd52009-07-22 20:33:47 +00001056 # Send a newline
1057 self.sendline()
1058 # Wait up to timeout seconds for some output from the child
1059 end_time = time.time() + timeout
1060 while time.time() < end_time:
1061 time.sleep(0.5)
1062 if self.read_nonblocking(timeout=0).strip():
1063 return True
1064 # No output -- report unresponsive
1065 return False
1066
1067
1068 def read_up_to_prompt(self, timeout=30.0, internal_timeout=None,
1069 print_func=None):
1070 """
1071 Read using read_nonblocking until the last non-empty line of the output
1072 matches the prompt regular expression set by set_prompt, or until
1073 timeout expires.
1074
1075 @brief: Read using read_nonblocking until the last non-empty line
1076 matches the prompt.
1077
1078 @param timeout: The duration (in seconds) to wait until a match is
1079 found
1080 @param internal_timeout: The timeout to pass to read_nonblocking
1081 @param print_func: A function to be used to print the data being
1082 read (should take a string parameter)
1083
1084 @return: A tuple containing True/False indicating whether the prompt
1085 was found, and the data read so far.
1086 """
1087 (match, output) = self.read_until_last_line_matches([self.prompt],
1088 timeout,
1089 internal_timeout,
1090 print_func)
1091 return (match is not None, output)
1092
1093
1094 def get_command_status_output(self, command, timeout=30.0,
1095 internal_timeout=None, print_func=None):
1096 """
1097 Send a command and return its exit status and output.
1098
1099 @param command: Command to send (must not contain newline characters)
1100 @param timeout: The duration (in seconds) to wait until a match is
1101 found
1102 @param internal_timeout: The timeout to pass to read_nonblocking
1103 @param print_func: A function to be used to print the data being read
1104 (should take a string parameter)
1105
1106 @return: A tuple (status, output) where status is the exit status or
1107 None if no exit status is available (e.g. timeout elapsed), and
1108 output is the output of command.
1109 """
1110 def remove_command_echo(str, cmd):
1111 if str and str.splitlines()[0] == cmd:
1112 str = "".join(str.splitlines(True)[1:])
1113 return str
1114
1115 def remove_last_nonempty_line(str):
1116 return "".join(str.rstrip().splitlines(True)[:-1])
1117
1118 # Print some debugging info
1119 logging.debug("Sending command: %s" % command)
1120
1121 # Read everything that's waiting to be read
lmr59b98db2009-10-05 19:11:21 +00001122 self.read_nonblocking(timeout=0)
lmrbbc9dd52009-07-22 20:33:47 +00001123
1124 # Send the command and get its output
1125 self.sendline(command)
1126 (match, output) = self.read_up_to_prompt(timeout, internal_timeout,
1127 print_func)
1128 # Remove the echoed command from the output
1129 output = remove_command_echo(output, command)
1130 # If the prompt was not found, return the output so far
1131 if not match:
1132 return (None, output)
1133 # Remove the final shell prompt from the output
1134 output = remove_last_nonempty_line(output)
1135
1136 # Send the 'echo ...' command to get the last exit status
1137 self.sendline(self.status_test_command)
1138 (match, status) = self.read_up_to_prompt(10.0, internal_timeout)
1139 if not match:
1140 return (None, output)
1141 status = remove_command_echo(status, self.status_test_command)
1142 status = remove_last_nonempty_line(status)
1143 # Get the first line consisting of digits only
1144 digit_lines = [l for l in status.splitlines() if l.strip().isdigit()]
1145 if not digit_lines:
1146 return (None, output)
1147 status = int(digit_lines[0].strip())
1148
1149 # Print some debugging info
1150 if status != 0:
1151 logging.debug("Command failed; status: %d, output:%s", status,
1152 kvm_utils.format_str_for_message(output.strip()))
1153
1154 return (status, output)
1155
1156
1157 def get_command_status(self, command, timeout=30.0, internal_timeout=None,
1158 print_func=None):
1159 """
1160 Send a command and return its exit status.
1161
1162 @param command: Command to send
1163 @param timeout: The duration (in seconds) to wait until a match is
1164 found
1165 @param internal_timeout: The timeout to pass to read_nonblocking
1166 @param print_func: A function to be used to print the data being read
1167 (should take a string parameter)
1168
1169 @return: Exit status or None if no exit status is available (e.g.
1170 timeout elapsed).
1171 """
1172 (status, output) = self.get_command_status_output(command, timeout,
1173 internal_timeout,
1174 print_func)
1175 return status
1176
1177
1178 def get_command_output(self, command, timeout=30.0, internal_timeout=None,
1179 print_func=None):
1180 """
1181 Send a command and return its output.
1182
1183 @param command: Command to send
1184 @param timeout: The duration (in seconds) to wait until a match is
1185 found
1186 @param internal_timeout: The timeout to pass to read_nonblocking
1187 @param print_func: A function to be used to print the data being read
1188 (should take a string parameter)
1189 """
1190 (status, output) = self.get_command_status_output(command, timeout,
1191 internal_timeout,
1192 print_func)
1193 return output