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