blob: 27bc7c13a82e961f287b3f42fb63823c34b4797a [file] [log] [blame]
lmr6f669ce2009-05-31 19:02:42 +00001import md5, thread, subprocess, time, string, random, socket, os, signal, pty
2import select, re, logging
lmr3f0b0cc2009-06-10 02:25:23 +00003from autotest_lib.client.bin import utils
4from autotest_lib.client.common_lib import error
lmr6f669ce2009-05-31 19:02:42 +00005
6"""
7KVM test utility functions.
8
9@copyright: 2008-2009 Red Hat Inc.
10"""
11
12
13def get_sub_dict(dict, name):
14 """
15 Return a "sub-dict" corresponding to a specific object.
16
17 Operate on a copy of dict: for each key that ends with the suffix
18 "_" + name, strip the suffix from the key, and set the value of
19 the stripped key to that of the key. Return the resulting dict.
20
21 @param name: Suffix of the key we want to set the value.
22 """
23 suffix = "_" + name
24 new_dict = dict.copy()
25 for key in dict.keys():
26 if key.endswith(suffix):
27 new_key = key.split(suffix)[0]
28 new_dict[new_key] = dict[key]
29 return new_dict
30
31
32def get_sub_dict_names(dict, keyword):
33 """
34 Return a list of "sub-dict" names that may be extracted with get_sub_dict.
35
36 This function may be modified to change the behavior of all functions that
37 deal with multiple objects defined in dicts (e.g. VMs, images, NICs).
38
39 @param keyword: A key in dict (e.g. "vms", "images", "nics").
40 """
41 names = dict.get(keyword)
42 if names:
43 return names.split()
44 else:
45 return []
46
47
48# Functions for working with the environment (a dict-like object)
49
50def is_vm(obj):
51 """
52 Tests whether a given object is a VM object.
53
54 @param obj: Python object (pretty much everything on python).
55 """
56 return obj.__class__.__name__ == "VM"
57
58
59def env_get_all_vms(env):
60 """
61 Return a list of all VM objects on a given environment.
62
63 @param env: Dictionary with environment items.
64 """
65 vms = []
66 for obj in env.values():
67 if is_vm(obj):
68 vms.append(obj)
69 return vms
70
71
72def env_get_vm(env, name):
73 """
74 Return a VM object by its name.
75
76 @param name: VM name.
77 """
78 return env.get("vm__%s" % name)
79
80
81def env_register_vm(env, name, vm):
82 """
83 Register a given VM in a given env.
84
85 @param env: Environment where we will register the VM.
86 @param name: VM name.
87 @param vm: VM object.
88 """
89 env["vm__%s" % name] = vm
90
91
92def env_unregister_vm(env, name):
93 """
94 Remove a given VM from a given env.
95
96 @param env: Environment where we will un-register the VM.
97 @param name: VM name.
98 """
99 del env["vm__%s" % name]
100
101
102# Utility functions for dealing with external processes
103
104def pid_exists(pid):
105 """
106 Return True if a given PID exists.
107
108 @param pid: Process ID number.
109 """
110 try:
111 os.kill(pid, 0)
112 return True
113 except:
114 return False
115
116
117def safe_kill(pid, signal):
118 """
119 Attempt to send a signal to a given process that may or may not exist.
120
121 @param signal: Signal number.
122 """
123 try:
124 os.kill(pid, signal)
125 return True
126 except:
127 return False
128
129
lmr3f0b0cc2009-06-10 02:25:23 +0000130def get_latest_kvm_release_tag(release_dir):
131 """
132 Fetches the latest release tag for KVM.
133
134 @param release_dir: KVM source forge download location.
135 """
136 try:
137 page_url = os.path.join(release_dir, "showfiles.php")
138 local_web_page = utils.unmap_url("/", page_url, "/tmp")
139 f = open(local_web_page, "r")
140 data = f.read()
141 f.close()
142 rx = re.compile("package_id=(\d+).*\>kvm\<", re.IGNORECASE)
143 matches = rx.findall(data)
144 package_id = matches[0]
145 #package_id = 209008
146 rx = re.compile("package_id=%s.*release_id=\d+\">(\d+)" % package_id,
147 re.IGNORECASE)
148 matches = rx.findall(data)
149 return matches[0] # the first match contains the latest release tag
150 except Exception, e:
151 message = "Could not fetch latest KVM release tag: %s" % str(e)
152 logging.error(message)
153 raise error.TestError(message)
154
155
156def get_git_branch(repository, branch, srcdir, commit=None, lbranch=None):
157 """
158 Retrieves a given git code repository.
159
160 @param repository: Git repository URL
161 """
162 logging.info("Fetching git [REP '%s' BRANCH '%s' TAG '%s'] -> %s",
163 repository, branch, commit, srcdir)
164 if not os.path.exists(srcdir):
165 os.makedirs(srcdir)
166 os.chdir(srcdir)
167
168 if os.path.exists(".git"):
169 utils.system("git reset --hard")
170 else:
171 utils.system("git init")
172
173 if not lbranch:
174 lbranch = branch
175
176 utils.system("git fetch -q -f -u -t %s %s:%s" %
177 (repository, branch, lbranch))
178 utils.system("git checkout %s" % lbranch)
179 if commit:
180 utils.system("git checkout %s" % commit)
181
182 h = utils.system_output('git log --pretty=format:"%H" -1')
183 desc = utils.system_output("git describe")
184 logging.info("Commit hash for %s is %s (%s)" % (repository, h.strip(),
185 desc))
186 return srcdir
187
188
189def unload_module(module_name):
190 """
191 Removes a module. Handles dependencies. If even then it's not possible
192 to remove one of the modules, it will trhow an error.CmdError exception.
193
194 @param module_name: Name of the module we want to remove.
195 """
196 l_raw = utils.system_output("/sbin/lsmod").splitlines()
197 lsmod = [x for x in l_raw if x.split()[0] == module_name]
198 if len(lsmod) > 0:
199 line_parts = lsmod[0].split()
200 if len(line_parts) == 4:
201 submodules = line_parts[3].split(",")
202 for submodule in submodules:
203 unload_module(submodule)
204 utils.system("/sbin/modprobe -r %s" % module_name)
205 logging.info("Module %s unloaded" % module_name)
206 else:
207 logging.info("Module %s is already unloaded" % module_name)
208
209
210def check_kvm_source_dir(source_dir):
211 """
212 Inspects the kvm source directory and verifies its disposition. In some
213 occasions build may be dependant on the source directory disposition.
214 The reason why the return codes are numbers is that we might have more
215 changes on the source directory layout, so it's not scalable to just use
216 strings like 'old_repo', 'new_repo' and such.
217
218 @param source_dir: Source code path that will be inspected.
219 """
220 os.chdir(source_dir)
221 has_qemu_dir = os.path.isdir('qemu')
222 has_kvm_dir = os.path.isdir('kvm')
223 if has_qemu_dir and not has_kvm_dir:
224 logging.debug("qemu directory detected, source dir layout 1")
225 return 1
226 if has_kvm_dir and not has_qemu_dir:
227 logging.debug("kvm directory detected, source dir layout 2")
228 return 2
229 else:
230 raise error.TestError("Unknown source dir layout, cannot proceed.")
231
232
lmr6f669ce2009-05-31 19:02:42 +0000233# The following are a class and functions used for SSH, SCP and Telnet
234# communication with guests.
235
236class kvm_spawn:
237 """
238 This class is used for spawning and controlling a child process.
239 """
240
241 def __init__(self, command, linesep="\n"):
242 """
243 Initialize the class and run command as a child process.
244
245 @param command: Command that will be run.
246 @param linesep: Line separator for the given platform.
247 """
248 self.exitstatus = None
249 self.linesep = linesep
250 (pid, fd) = pty.fork()
251 if pid == 0:
252 os.execv("/bin/sh", ["/bin/sh", "-c", command])
253 else:
254 self.pid = pid
255 self.fd = fd
256
257
258 def set_linesep(self, linesep):
259 """
260 Sets the line separator string (usually "\\n").
261
262 @param linesep: Line separator character.
263 """
264 self.linesep = linesep
265
266
267 def is_responsive(self, timeout=5.0):
268 """
269 Return True if the session is responsive.
270
271 Send a newline to the child process (e.g. SSH or Telnet) and read some
272 output using read_nonblocking.
273 If all is OK, some output should be available (e.g. the shell prompt).
274 In that case return True. Otherwise return False.
275
276 @param timeout: Timeout that will happen before we consider the
277 process unresponsive
278 """
279 self.read_nonblocking(timeout=0.1)
280 self.sendline()
281 output = self.read_nonblocking(timeout=timeout)
282 if output.strip():
283 return True
284 return False
285
286
287 def poll(self):
288 """
289 If the process exited, return its exit status. Otherwise return None.
290 The exit status is stored for use in subsequent calls.
291 """
292 if self.exitstatus != None:
293 return self.exitstatus
294 pid, status = os.waitpid(self.pid, os.WNOHANG)
295 if pid:
296 self.exitstatus = os.WEXITSTATUS(status)
297 return self.exitstatus
298 else:
299 return None
300
301
302 def close(self):
303 """
304 Close the session (close the process filedescriptors and kills the
305 process ID), and return the exit status.
306 """
307 try:
308 os.close(self.fd)
309 os.kill(self.pid, signal.SIGTERM)
310 except OSError:
311 pass
312 return self.poll()
313
314
315 def sendline(self, str=""):
316 """
317 Sends a string followed by a line separator to the child process.
318
319 @param str: String that will be sent to the child process.
320 """
321 try:
322 os.write(self.fd, str + self.linesep)
323 except OSError:
324 pass
325
326
327 def read_nonblocking(self, timeout=1.0):
328 """
329 Read from child until there is nothing to read for timeout seconds.
330
331 @param timeout: Time (seconds) of wait before we give up reading from
332 the child process.
333 """
334 data = ""
335 while True:
336 r, w, x = select.select([self.fd], [], [], timeout)
337 if self.fd in r:
338 try:
339 data += os.read(self.fd, 1024)
340 except OSError:
341 return data
342 else:
343 return data
344
345
346 def match_patterns(self, str, patterns):
347 """
348 Match str against a list of patterns.
349
350 Return the index of the first pattern that matches a substring of str.
351 None and empty strings in patterns are ignored.
352 If no match is found, return None.
353
354 @param patterns: List of strings (regular expression patterns).
355 """
356 for i in range(len(patterns)):
357 if not patterns[i]:
358 continue
359 exp = re.compile(patterns[i])
360 if exp.search(str):
361 return i
362
363
364 def read_until_output_matches(self, patterns, filter=lambda(x):x,
365 timeout=30.0, internal_timeout=1.0,
366 print_func=None):
367 """
368 Read using read_nonblocking until a match is found using match_patterns,
369 or until timeout expires. Before attempting to search for a match, the
370 data is filtered using the filter function provided.
371
372 @brief: Read from child using read_nonblocking until a pattern
373 matches.
374 @param patterns: List of strings (regular expression patterns)
375 @param filter: Function to apply to the data read from the child before
376 attempting to match it against the patterns (should take and
377 return a string)
378 @param timeout: The duration (in seconds) to wait until a match is
379 found
380 @param internal_timeout: The timeout to pass to read_nonblocking
381 @param print_func: A function to be used to print the data being read
382 (should take a string parameter)
383 @return: Tuple containing the match index (or None if no match was
384 found) and the data read so far.
385 """
386 match = None
387 data = ""
388
389 end_time = time.time() + timeout
390 while time.time() < end_time:
391 # Read data from child
392 newdata = self.read_nonblocking(internal_timeout)
393 # Print it if necessary
394 if print_func and newdata:
395 map(print_func, newdata.splitlines())
396 data += newdata
397
398 done = False
399 # Look for patterns
400 match = self.match_patterns(filter(data), patterns)
401 if match != None:
402 done = True
403 # Check if child has died
404 if self.poll() != None:
lmr3f0b0cc2009-06-10 02:25:23 +0000405 logging.debug("Process terminated with status %d", self.poll())
lmr6f669ce2009-05-31 19:02:42 +0000406 done = True
407 # Are we done?
408 if done: break
409
410 # Print some debugging info
411 if match == None and self.poll() != 0:
lmr3f0b0cc2009-06-10 02:25:23 +0000412 logging.debug("Timeout elapsed or process terminated. Output: %s",
lmr6f669ce2009-05-31 19:02:42 +0000413 format_str_for_message(data.strip()))
414
415 return (match, data)
416
417
418 def get_last_word(self, str):
419 """
420 Return the last word in str.
421
422 @param str: String that will be analyzed.
423 """
424 if str:
425 return str.split()[-1]
426 else:
427 return ""
428
429
430 def get_last_line(self, str):
431 """
432 Return the last non-empty line in str.
433
434 @param str: String that will be analyzed.
435 """
436 last_line = ""
437 for line in str.splitlines():
438 if line != "":
439 last_line = line
440 return last_line
441
442
443 def read_until_last_word_matches(self, patterns, timeout=30.0,
444 internal_timeout=1.0, print_func=None):
445 """
446 Read using read_nonblocking until the last word of the output matches
447 one of the patterns (using match_patterns), or until timeout expires.
448
449 @param patterns: A list of strings (regular expression patterns)
450 @param timeout: The duration (in seconds) to wait until a match is
451 found
452 @param internal_timeout: The timeout to pass to read_nonblocking
453 @param print_func: A function to be used to print the data being read
454 (should take a string parameter)
455 @return: A tuple containing the match index (or None if no match was
456 found) and the data read so far.
457 """
458 return self.read_until_output_matches(patterns, self.get_last_word,
459 timeout, internal_timeout,
460 print_func)
461
462
463 def read_until_last_line_matches(self, patterns, timeout=30.0,
464 internal_timeout=1.0, print_func=None):
465 """
466 Read using read_nonblocking until the last non-empty line of the output
467 matches one of the patterns (using match_patterns), or until timeout
468 expires. Return a tuple containing the match index (or None if no match
469 was found) and the data read so far.
470
471 @brief: Read using read_nonblocking until the last non-empty line
472 matches a pattern.
473
474 @param patterns: A list of strings (regular expression patterns)
475 @param timeout: The duration (in seconds) to wait until a match is
476 found
477 @param internal_timeout: The timeout to pass to read_nonblocking
478 @param print_func: A function to be used to print the data being read
479 (should take a string parameter)
480 """
481 return self.read_until_output_matches(patterns, self.get_last_line,
482 timeout, internal_timeout,
483 print_func)
484
485
486 def set_prompt(self, prompt):
487 """
488 Set the prompt attribute for later use by read_up_to_prompt.
489
490 @param: String that describes the prompt contents.
491 """
492 self.prompt = prompt
493
494
495 def read_up_to_prompt(self, timeout=30.0, internal_timeout=1.0,
496 print_func=None):
497 """
498 Read using read_nonblocking until the last non-empty line of the output
499 matches the prompt regular expression set by set_prompt, or until
500 timeout expires.
501
502 @brief: Read using read_nonblocking until the last non-empty line
503 matches the prompt.
504
505 @param timeout: The duration (in seconds) to wait until a match is
506 found
507 @param internal_timeout: The timeout to pass to read_nonblocking
508 @param print_func: A function to be used to print the data being
509 read (should take a string parameter)
510
511 @return: A tuple containing True/False indicating whether the prompt
512 was found, and the data read so far.
513 """
514 (match, output) = self.read_until_last_line_matches([self.prompt],
515 timeout,
516 internal_timeout,
517 print_func)
518 if match == None:
519 return (False, output)
520 else:
521 return (True, output)
522
523
524 def set_status_test_command(self, status_test_command):
525 """
526 Set the command to be sent in order to get the last exit status.
527
528 @param status_test_command: Command that will be sent to get the last
529 exit status.
530 """
531 self.status_test_command = status_test_command
532
533
534 def get_command_status_output(self, command, timeout=30.0,
535 internal_timeout=1.0, print_func=None):
536 """
537 Send a command and return its exit status and output.
538
539 @param command: Command to send
540 @param timeout: The duration (in seconds) to wait until a match is
541 found
542 @param internal_timeout: The timeout to pass to read_nonblocking
543 @param print_func: A function to be used to print the data being read
544 (should take a string parameter)
545
546 @return: A tuple (status, output) where status is the exit status or
547 None if no exit status is available (e.g. timeout elapsed), and
548 output is the output of command.
549 """
550 # Print some debugging info
551 logging.debug("Sending command: %s" % command)
552
553 # Read everything that's waiting to be read
554 self.read_nonblocking(0.1)
555
556 # Send the command and get its output
557 self.sendline(command)
558 (match, output) = self.read_up_to_prompt(timeout, internal_timeout,
559 print_func)
560 if not match:
561 return (None, "\n".join(output.splitlines()[1:]))
562 output = "\n".join(output.splitlines()[1:-1])
563
564 # Send the 'echo ...' command to get the last exit status
565 self.sendline(self.status_test_command)
566 (match, status) = self.read_up_to_prompt(10.0, internal_timeout)
567 if not match:
568 return (None, output)
569 status = int("\n".join(status.splitlines()[1:-1]).strip())
570
571 # Print some debugging info
572 if status != 0:
573 logging.debug("Command failed; status: %d, output:" % status \
574 + format_str_for_message(output.strip()))
575
576 return (status, output)
577
578
579 def get_command_status(self, command, timeout=30.0, internal_timeout=1.0,
580 print_func=None):
581 """
582 Send a command and return its exit status.
583
584 @param command: Command to send
585 @param timeout: The duration (in seconds) to wait until a match is
586 found
587 @param internal_timeout: The timeout to pass to read_nonblocking
588 @param print_func: A function to be used to print the data being read
589 (should take a string parameter)
590
591 @return: Exit status or None if no exit status is available (e.g.
592 timeout elapsed).
593 """
594 (status, output) = self.get_command_status_output(command, timeout,
595 internal_timeout,
596 print_func)
597 return status
598
599
600 def get_command_output(self, command, timeout=30.0, internal_timeout=1.0,
601 print_func=None):
602 """
603 Send a command and return its output.
604
605 @param command: Command to send
606 @param timeout: The duration (in seconds) to wait until a match is
607 found
608 @param internal_timeout: The timeout to pass to read_nonblocking
609 @param print_func: A function to be used to print the data being read
610 (should take a string parameter)
611 """
612 (status, output) = self.get_command_status_output(command, timeout,
613 internal_timeout,
614 print_func)
615 return output
616
617
618def remote_login(command, password, prompt, linesep="\n", timeout=10):
619 """
620 Log into a remote host (guest) using SSH or Telnet. Run the given command
621 using kvm_spawn and provide answers to the questions asked. If timeout
622 expires while waiting for output from the child (e.g. a password prompt
623 or a shell prompt) -- fail.
624
625 @brief: Log into a remote host (guest) using SSH or Telnet.
626
627 @param command: The command to execute (e.g. "ssh root@localhost")
628 @param password: The password to send in reply to a password prompt
629 @param prompt: The shell prompt that indicates a successful login
630 @param linesep: The line separator to send instead of "\\n"
631 (sometimes "\\r\\n" is required)
632 @param timeout: The maximal time duration (in seconds) to wait for each
633 step of the login procedure (i.e. the "Are you sure" prompt, the
634 password prompt, the shell prompt, etc)
635
636 @return Return the kvm_spawn object on success and None on failure.
637 """
638 sub = kvm_spawn(command, linesep)
639 sub.set_prompt(prompt)
640
641 password_prompt_count = 0
642
643 logging.debug("Trying to login...")
644
645 while True:
646 (match, text) = sub.read_until_last_line_matches(
647 ["[Aa]re you sure", "[Pp]assword:", "[Ll]ogin:",
648 "[Cc]onnection.*closed", prompt],
649 timeout=timeout, internal_timeout=0.5)
650 if match == 0: # "Are you sure you want to continue connecting"
651 logging.debug("Got 'Are you sure...'; sending 'yes'")
652 sub.sendline("yes")
653 continue
654 elif match == 1: # "password:"
655 if password_prompt_count == 0:
656 logging.debug("Got password prompt; sending '%s'" % password)
657 sub.sendline(password)
658 password_prompt_count += 1
659 continue
660 else:
661 logging.debug("Got password prompt again")
662 sub.close()
663 return None
664 elif match == 2: # "login:"
665 logging.debug("Got unexpected login prompt")
666 sub.close()
667 return None
668 elif match == 3: # "Connection closed"
669 logging.debug("Got 'Connection closed'")
670 sub.close()
671 return None
672 elif match == 4: # prompt
673 logging.debug("Got shell prompt -- logged in")
674 return sub
675 else: # match == None
676 logging.debug("Timeout or process terminated")
677 sub.close()
678 return None
679
680
681def remote_scp(command, password, timeout=300, login_timeout=10):
682 """
683 Run the given command using kvm_spawn and provide answers to the questions
684 asked. If timeout expires while waiting for the transfer to complete ,
685 fail. If login_timeout expires while waiting for output from the child
686 (e.g. a password prompt), fail.
687
688 @brief: Transfer files using SCP, given a command line.
689
690 @param command: The command to execute
691 (e.g. "scp -r foobar root@localhost:/tmp/").
692 @param password: The password to send in reply to a password prompt.
693 @param timeout: The time duration (in seconds) to wait for the transfer
694 to complete.
695 @param login_timeout: The maximal time duration (in seconds) to wait for
696 each step of the login procedure (i.e. the "Are you sure" prompt or the
697 password prompt)
698
699 @return: True if the transfer succeeds and False on failure.
700 """
701 sub = kvm_spawn(command)
702
703 password_prompt_count = 0
704 _timeout = login_timeout
705
706 logging.debug("Trying to login...")
707
708 while True:
709 (match, text) = sub.read_until_last_line_matches(
710 ["[Aa]re you sure", "[Pp]assword:", "lost connection"],
711 timeout=_timeout, internal_timeout=0.5)
712 if match == 0: # "Are you sure you want to continue connecting"
713 logging.debug("Got 'Are you sure...'; sending 'yes'")
714 sub.sendline("yes")
715 continue
716 elif match == 1: # "password:"
717 if password_prompt_count == 0:
718 logging.debug("Got password prompt; sending '%s'" % password)
719 sub.sendline(password)
720 password_prompt_count += 1
721 _timeout = timeout
722 continue
723 else:
724 logging.debug("Got password prompt again")
725 sub.close()
726 return False
727 elif match == 2: # "lost connection"
728 logging.debug("Got 'lost connection'")
729 sub.close()
730 return False
731 else: # match == None
732 logging.debug("Timeout or process terminated")
733 sub.close()
734 return sub.poll() == 0
735
736
737def scp_to_remote(host, port, username, password, local_path, remote_path,
738 timeout=300):
739 """
740 Copy files to a remote host (guest).
741
742 @param host: Hostname of the guest
743 @param username: User that will be used to copy the files
744 @param password: Host's password
745 @param local_path: Path on the local machine where we are copying from
746 @param remote_path: Path on the remote machine where we are copying to
747 @param timeout: Time in seconds that we will wait before giving up to
748 copy the files.
749
750 @return: True on success and False on failure.
751 """
752 command = "scp -o UserKnownHostsFile=/dev/null -r -P %s %s %s@%s:%s" % \
753 (port, local_path, username, host, remote_path)
754 return remote_scp(command, password, timeout)
755
756
757def scp_from_remote(host, port, username, password, remote_path, local_path,
758 timeout=300):
759 """
760 Copy files from a remote host (guest).
761
762 @param host: Hostname of the guest
763 @param username: User that will be used to copy the files
764 @param password: Host's password
765 @param local_path: Path on the local machine where we are copying from
766 @param remote_path: Path on the remote machine where we are copying to
767 @param timeout: Time in seconds that we will wait before giving up to copy
768 the files.
769
770 @return: True on success and False on failure.
771 """
772 command = "scp -o UserKnownHostsFile=/dev/null -r -P %s %s@%s:%s %s" % \
773 (port, username, host, remote_path, local_path)
774 return remote_scp(command, password, timeout)
775
776
777def ssh(host, port, username, password, prompt, timeout=10):
778 """
779 Log into a remote host (guest) using SSH.
780
781 @param host: Hostname of the guest
782 @param username: User that will be used to log into the host.
783 @param password: Host's password
784 @timeout: Time in seconds that we will wait before giving up on logging
785 into the host.
786
787 @return: kvm_spawn object on success and None on failure.
788 """
789 command = "ssh -o UserKnownHostsFile=/dev/null -p %s %s@%s" % \
790 (port, username, host)
791 return remote_login(command, password, prompt, "\n", timeout)
792
793
794def telnet(host, port, username, password, prompt, timeout=10):
795 """
796 Log into a remote host (guest) using Telnet.
797
798 @param host: Hostname of the guest
799 @param username: User that will be used to log into the host.
800 @param password: Host's password
801 @timeout: Time in seconds that we will wait before giving up on logging
802 into the host.
803
804 @return: kvm_spawn object on success and None on failure.
805 """
806 command = "telnet -l %s %s %s" % (username, host, port)
807 return remote_login(command, password, prompt, "\r\n", timeout)
808
809
810# The following are functions used for running commands in the background.
811
812def track_process(sub, status_output=None, term_func=None, stdout_func=None,
813 prefix=""):
814 """
815 Read lines from the stdout pipe of the subprocess. Pass each line to
816 stdout_func prefixed by prefix. Place the lines in status_output[1].
817 When the subprocess exits, call term_func. Place the exit status in
818 status_output[0].
819
820 @brief Track a subprocess and report its output and termination.
821
822 @param sub: An object returned by subprocess.Popen
823 @param status_output: A list in which the exit status and output are to be
824 stored.
825 @param term_func: A function to call when the process terminates
826 (should take no parameters)
827 @param stdout_func: A function to call with each line of output from the
828 subprocess (should take a string parameter)
829
830 @param prefix -- a string to pre-pend to each line of the output, before
831 passing it to stdout_func
832 """
833 while True:
834 # Read a single line from stdout
835 text = sub.stdout.readline()
836 # If the subprocess exited...
837 if text == "":
838 # Get exit code
839 status = sub.wait()
840 # Report it
841 if status_output:
842 status_output[0] = status
843 # Call term_func
844 if term_func:
845 term_func()
846 return
847 # Report the text
848 if status_output:
849 status_output[1] += text
850 # Call stdout_func with the returned text
851 if stdout_func:
852 text = prefix + text.strip()
853 # We need to sanitize the text before passing it to the logging
854 # system
855 text = text.decode('utf-8', 'replace')
856 stdout_func(text)
857
858
859def run_bg(command, term_func=None, stdout_func=None, prefix="", timeout=1.0):
860 """
861 Run command as a subprocess. Call stdout_func with each line of output from
862 the subprocess (prefixed by prefix). Call term_func when the subprocess
863 terminates. If timeout expires and the subprocess is still running, return.
864
865 @brief: Run a subprocess in the background and collect its output and
866 exit status.
867
868 @param command: The shell command to execute
869 @param term_func: A function to call when the process terminates
870 (should take no parameters)
871 @param stdout_func: A function to call with each line of output from
872 the subprocess (should take a string parameter)
873 @param prefix: A string to pre-pend to each line of the output, before
874 passing it to stdout_func
875 @param timeout: Time duration (in seconds) to wait for the subprocess to
876 terminate before returning
877
878 @return: A 3-tuple containing the exit status (None if the subprocess is
879 still running), the PID of the subprocess (None if the subprocess
880 terminated), and the output collected so far.
881 """
882 # Start the process
883 sub = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE,
884 stderr=subprocess.STDOUT)
885 # Start the tracking thread
886 status_output = [None, ""]
887 thread.start_new_thread(track_process, (sub, status_output, term_func,
888 stdout_func, prefix))
889 # Wait up to timeout secs for the process to exit
890 end_time = time.time() + timeout
891 while time.time() < end_time:
892 # If the process exited, return
893 if status_output[0] != None:
894 return (status_output[0], None, status_output[1])
895 # Otherwise, sleep for a while
896 time.sleep(0.1)
897 # Report the PID and the output collected so far
898 return (None, sub.pid, status_output[1])
899
900
901# The following are utility functions related to ports.
902
903def is_sshd_running(host, port, timeout=10.0):
904 """
905 Connect to the given host and port and wait for output.
906 Return True if the given host and port are responsive.
907
908 @param host: Host's hostname
909 @param port: Host's port
910 @param timeout: Time (seconds) before we giving up on checking the SSH
911 daemon.
912
913 @return: If output is available, return True. If timeout expires and no
914 output was available, return False.
915 """
916 try:
917 # Try to connect
918 #s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
919 s = socket.socket()
920 s.connect((host, port))
921 except socket.error:
922 # Can't connect -- return False
923 s.close()
924 logging.debug("Could not connect")
925 return False
926 s.setblocking(False)
927 # Wait up to 'timeout' seconds
928 end_time = time.time() + timeout
929 while time.time() < end_time:
930 try:
931 time.sleep(0.1)
932 # Try to receive some text
933 str = s.recv(1024)
934 if len(str) > 0:
935 s.shutdown(socket.SHUT_RDWR)
936 s.close()
937 logging.debug("Success! got string %r" % str)
938 return True
939 except socket.error:
940 # No text was available; try again
941 pass
942 # Timeout elapsed and no text was received
943 s.shutdown(socket.SHUT_RDWR)
944 s.close()
945 logging.debug("Timeout")
946 return False
947
948
949def is_port_free(port):
950 """
951 Return True if the given port is available for use.
952
953 @param port: Port number
954 """
955 try:
956 s = socket.socket()
957 #s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
958 s.bind(("localhost", port))
959 free = True
960 except socket.error:
961 free = False
962 s.close()
963 return free
964
965
966def find_free_port(start_port, end_port):
967 """
968 Return a free port in the range [start_port, end_port).
969
970 @param start_port: First port that will be checked.
971 @param end_port: Port immediately after the last one that will be checked.
972 """
973 for i in range(start_port, end_port):
974 if is_port_free(i):
975 return i
976 return None
977
978
979def find_free_ports(start_port, end_port, count):
980 """
981 Return count free ports in the range [start_port, end_port).
982
983 @count: Initial number of ports known to be free in the range.
984 @param start_port: First port that will be checked.
985 @param end_port: Port immediately after the last one that will be checked.
986 """
987 ports = []
988 i = start_port
989 while i < end_port and count > 0:
990 if is_port_free(i):
991 ports.append(i)
992 count -= 1
993 i += 1
994 return ports
995
996
997# The following are miscellaneous utility functions.
998
999def generate_random_string(length):
1000 """
1001 Return a random string using alphanumeric characters.
1002
1003 @length: length of the string that will be generated.
1004 """
1005 str = ""
1006 chars = string.letters + string.digits
1007 while length > 0:
1008 str += random.choice(chars)
1009 length -= 1
1010 return str
1011
1012
1013def format_str_for_message(str):
1014 """
1015 Format str so that it can be appended to a message.
1016 If str consists of one line, prefix it with a space.
1017 If str consists of multiple lines, prefix it with a newline.
1018
1019 @param str: string that will be formatted.
1020 """
1021 num_lines = len(str.splitlines())
1022 if num_lines == 0:
1023 return ""
1024 elif num_lines == 1:
1025 return " " + str
1026 else:
1027 return "\n" + str
1028
1029
1030def wait_for(func, timeout, first=0.0, step=1.0, text=None):
1031 """
1032 If func() evaluates to True before timeout expires, return the
1033 value of func(). Otherwise return None.
1034
1035 @brief: Wait until func() evaluates to True.
1036
1037 @param timeout: Timeout in seconds
1038 @param first: Time to sleep before first attempt
1039 @param steps: Time to sleep between attempts in seconds
1040 @param text: Text to print while waiting, for debug purposes
1041 """
1042 start_time = time.time()
1043 end_time = time.time() + timeout
1044
1045 time.sleep(first)
1046
1047 while time.time() < end_time:
1048 if text:
1049 logging.debug("%s (%f secs)" % (text, time.time() - start_time))
1050
1051 output = func()
1052 if output:
1053 return output
1054
1055 time.sleep(step)
1056
1057 logging.debug("Timeout elapsed")
1058 return None
1059
1060
1061def md5sum_file(filename, size=None):
1062 """
1063 Calculate the md5sum of filename.
1064 If size is not None, limit to first size bytes.
1065 Throw exception if something is wrong with filename.
1066 Can be also implemented with bash one-liner (assuming size%1024==0):
1067 dd if=filename bs=1024 count=size/1024 | md5sum -
1068
1069 @param filename: Path of the file that will have its md5sum calculated.
1070 @param returns: md5sum of the file.
1071 """
1072 chunksize = 4096
1073 fsize = os.path.getsize(filename)
1074 if not size or size>fsize:
1075 size = fsize
1076 f = open(filename, 'rb')
1077 o = md5.new()
1078 while size > 0:
1079 if chunksize > size:
1080 chunksize = size
1081 data = f.read(chunksize)
1082 if len(data) == 0:
1083 logging.debug("Nothing left to read but size=%d" % size)
1084 break
1085 o.update(data)
1086 size -= len(data)
1087 f.close()
1088 return o.hexdigest()