1import os, time, commands, re, logging, glob, threading, shutil
2from autotest_lib.client.bin import utils
3from autotest_lib.client.common_lib import error
4import aexpect, virt_utils, kvm_monitor, ppm_utils, virt_test_setup
5import virt_vm, kvm_vm
6try:
7    import PIL.Image
8except ImportError:
9    logging.warning('No python imaging library installed. PPM image '
10                    'conversion to JPEG disabled. In order to enable it, '
11                    'please install python-imaging or the equivalent for your '
12                    'distro.')
13
14
15_screendump_thread = None
16_screendump_thread_termination_event = None
17
18
19def preprocess_image(test, params):
20    """
21    Preprocess a single QEMU image according to the instructions in params.
22
23    @param test: Autotest test object.
24    @param params: A dict containing image preprocessing parameters.
25    @note: Currently this function just creates an image if requested.
26    """
27    image_filename = virt_vm.get_image_filename(params, test.bindir)
28
29    create_image = False
30
31    if params.get("force_create_image") == "yes":
32        logging.debug("Param 'force_create_image' specified, creating image")
33        create_image = True
34    elif (params.get("create_image") == "yes" and not
35          os.path.exists(image_filename)):
36        create_image = True
37
38    if create_image and not virt_vm.create_image(params, test.bindir):
39        raise error.TestError("Could not create image")
40
41
42def preprocess_vm(test, params, env, name):
43    """
44    Preprocess a single VM object according to the instructions in params.
45    Start the VM if requested and get a screendump.
46
47    @param test: An Autotest test object.
48    @param params: A dict containing VM preprocessing parameters.
49    @param env: The environment (a dict-like object).
50    @param name: The name of the VM object.
51    """
52    logging.debug("Preprocessing VM '%s'", name)
53    vm = env.get_vm(name)
54    if not vm:
55        logging.debug("VM object for '%s' does not exist, creating it", name)
56        vm_type = params.get('vm_type')
57        if vm_type == 'kvm':
58            vm = kvm_vm.VM(name, params, test.bindir, env.get("address_cache"))
59        env.register_vm(name, vm)
60
61    start_vm = False
62
63    if params.get("restart_vm") == "yes":
64        logging.debug("Param 'restart_vm' specified, (re)starting VM")
65        start_vm = True
66    elif params.get("migration_mode"):
67        logging.debug("Param 'migration_mode' specified, starting VM in "
68                      "incoming migration mode")
69        start_vm = True
70    elif params.get("start_vm") == "yes":
71        if not vm.is_alive():
72            logging.debug("VM is not alive, starting it")
73            start_vm = True
74        if vm.needs_restart(name=name, params=params, basedir=test.bindir):
75            logging.debug("Current VM specs differ from requested one; "
76                          "restarting it")
77            start_vm = True
78
79    if start_vm:
80        # Start the VM (or restart it if it's already up)
81        vm.create(name, params, test.bindir,
82                  migration_mode=params.get("migration_mode"))
83    else:
84        # Don't start the VM, just update its params
85        vm.params = params
86
87    scrdump_filename = os.path.join(test.debugdir, "pre_%s.ppm" % name)
88    try:
89        if vm.monitor and params.get("take_regular_screendumps") == "yes":
90            vm.monitor.screendump(scrdump_filename, debug=False)
91    except kvm_monitor.MonitorError, e:
92        logging.warning(e)
93
94
95def postprocess_image(test, params):
96    """
97    Postprocess a single QEMU image according to the instructions in params.
98
99    @param test: An Autotest test object.
100    @param params: A dict containing image postprocessing parameters.
101    """
102    if params.get("check_image") == "yes":
103        virt_vm.check_image(params, test.bindir)
104    if params.get("remove_image") == "yes":
105        virt_vm.remove_image(params, test.bindir)
106
107
108def postprocess_vm(test, params, env, name):
109    """
110    Postprocess a single VM object according to the instructions in params.
111    Kill the VM if requested and get a screendump.
112
113    @param test: An Autotest test object.
114    @param params: A dict containing VM postprocessing parameters.
115    @param env: The environment (a dict-like object).
116    @param name: The name of the VM object.
117    """
118    logging.debug("Postprocessing VM '%s'" % name)
119    vm = env.get_vm(name)
120    if not vm:
121        return
122
123    scrdump_filename = os.path.join(test.debugdir, "post_%s.ppm" % name)
124    try:
125        if vm.monitor and params.get("take_regular_screenshots") == "yes":
126            vm.monitor.screendump(scrdump_filename, debug=False)
127    except kvm_monitor.MonitorError, e:
128        logging.warning(e)
129
130    if params.get("kill_vm") == "yes":
131        kill_vm_timeout = float(params.get("kill_vm_timeout", 0))
132        if kill_vm_timeout:
133            logging.debug("Param 'kill_vm' specified, waiting for VM to shut "
134                          "down before killing it")
135            virt_utils.wait_for(vm.is_dead, kill_vm_timeout, 0, 1)
136        else:
137            logging.debug("Param 'kill_vm' specified, killing VM")
138        vm.destroy(gracefully = params.get("kill_vm_gracefully") == "yes")
139
140
141def process_command(test, params, env, command, command_timeout,
142                    command_noncritical):
143    """
144    Pre- or post- custom commands to be executed before/after a test is run
145
146    @param test: An Autotest test object.
147    @param params: A dict containing all VM and image parameters.
148    @param env: The environment (a dict-like object).
149    @param command: Command to be run.
150    @param command_timeout: Timeout for command execution.
151    @param command_noncritical: If True test will not fail if command fails.
152    """
153    # Export environment vars
154    for k in params:
155        os.putenv("KVM_TEST_%s" % k, str(params[k]))
156    # Execute commands
157    try:
158        utils.system("cd %s; %s" % (test.bindir, command))
159    except error.CmdError, e:
160        if command_noncritical:
161            logging.warning(e)
162        else:
163            raise
164
165
166def process(test, params, env, image_func, vm_func, vm_first=False):
167    """
168    Pre- or post-process VMs and images according to the instructions in params.
169    Call image_func for each image listed in params and vm_func for each VM.
170
171    @param test: An Autotest test object.
172    @param params: A dict containing all VM and image parameters.
173    @param env: The environment (a dict-like object).
174    @param image_func: A function to call for each image.
175    @param vm_func: A function to call for each VM.
176    """
177    # Get list of VMs specified for this test
178    for vm_name in params.objects("vms"):
179        vm_params = params.object_params(vm_name)
180        if not vm_first:
181            # Get list of images specified for this VM
182            for image_name in vm_params.objects("images"):
183                image_params = vm_params.object_params(image_name)
184                # Call image_func for each image
185                image_func(test, image_params)
186            # Call vm_func for each vm
187            vm_func(test, vm_params, env, vm_name)
188        else:
189            vm_func(test, vm_params, env, vm_name)
190            for image_name in vm_params.objects("images"):
191                image_params = vm_params.object_params(image_name)
192                image_func(test, image_params)
193
194
195
196@error.context_aware
197def preprocess(test, params, env):
198    """
199    Preprocess all VMs and images according to the instructions in params.
200    Also, collect some host information, such as the KVM version.
201
202    @param test: An Autotest test object.
203    @param params: A dict containing all VM and image parameters.
204    @param env: The environment (a dict-like object).
205    """
206    error.context("preprocessing")
207
208    # Start tcpdump if it isn't already running
209    if "address_cache" not in env:
210        env["address_cache"] = {}
211    if "tcpdump" in env and not env["tcpdump"].is_alive():
212        env["tcpdump"].close()
213        del env["tcpdump"]
214    if "tcpdump" not in env and params.get("run_tcpdump", "yes") == "yes":
215        cmd = "%s -npvi any 'dst port 68'" % virt_utils.find_command("tcpdump")
216        logging.debug("Starting tcpdump '%s'", cmd)
217        env["tcpdump"] = aexpect.Tail(
218            command=cmd,
219            output_func=_update_address_cache,
220            output_params=(env["address_cache"],))
221        if virt_utils.wait_for(lambda: not env["tcpdump"].is_alive(),
222                              0.1, 0.1, 1.0):
223            logging.warning("Could not start tcpdump")
224            logging.warning("Status: %s" % env["tcpdump"].get_status())
225            logging.warning("Output:" + virt_utils.format_str_for_message(
226                env["tcpdump"].get_output()))
227
228    # Destroy and remove VMs that are no longer needed in the environment
229    requested_vms = params.objects("vms")
230    for key in env.keys():
231        vm = env[key]
232        if not virt_utils.is_vm(vm):
233            continue
234        if not vm.name in requested_vms:
235            logging.debug("VM '%s' found in environment but not required for "
236                          "test, destroying it" % vm.name)
237            vm.destroy()
238            del env[key]
239
240    # Get the KVM kernel module version and write it as a keyval
241    if os.path.exists("/dev/kvm"):
242        try:
243            kvm_version = open("/sys/module/kvm/version").read().strip()
244        except:
245            kvm_version = os.uname()[2]
246    else:
247        kvm_version = "Unknown"
248        logging.debug("KVM module not loaded")
249    logging.debug("KVM version: %s" % kvm_version)
250    test.write_test_keyval({"kvm_version": kvm_version})
251
252    # Get the KVM userspace version and write it as a keyval
253    qemu_path = virt_utils.get_path(test.bindir, params.get("qemu_binary",
254                                                           "qemu"))
255    version_line = commands.getoutput("%s -help | head -n 1" % qemu_path)
256    matches = re.findall("[Vv]ersion .*?,", version_line)
257    if matches:
258        kvm_userspace_version = " ".join(matches[0].split()[1:]).strip(",")
259    else:
260        kvm_userspace_version = "Unknown"
261    logging.debug("KVM userspace version: %s" % kvm_userspace_version)
262    test.write_test_keyval({"kvm_userspace_version": kvm_userspace_version})
263
264    if params.get("setup_hugepages") == "yes":
265        h = virt_test_setup.HugePageConfig(params)
266        h.setup()
267
268    # Execute any pre_commands
269    if params.get("pre_command"):
270        process_command(test, params, env, params.get("pre_command"),
271                        int(params.get("pre_command_timeout", "600")),
272                        params.get("pre_command_noncritical") == "yes")
273
274    # Preprocess all VMs and images
275    process(test, params, env, preprocess_image, preprocess_vm)
276
277    # Start the screendump thread
278    if params.get("take_regular_screendumps") == "yes":
279        logging.debug("Starting screendump thread")
280        global _screendump_thread, _screendump_thread_termination_event
281        _screendump_thread_termination_event = threading.Event()
282        _screendump_thread = threading.Thread(target=_take_screendumps,
283                                              args=(test, params, env))
284        _screendump_thread.start()
285
286
287@error.context_aware
288def postprocess(test, params, env):
289    """
290    Postprocess all VMs and images according to the instructions in params.
291
292    @param test: An Autotest test object.
293    @param params: Dict containing all VM and image parameters.
294    @param env: The environment (a dict-like object).
295    """
296    error.context("postprocessing")
297
298    # Postprocess all VMs and images
299    process(test, params, env, postprocess_image, postprocess_vm, vm_first=True)
300
301    # Terminate the screendump thread
302    global _screendump_thread, _screendump_thread_termination_event
303    if _screendump_thread:
304        logging.debug("Terminating screendump thread")
305        _screendump_thread_termination_event.set()
306        _screendump_thread.join(10)
307        _screendump_thread = None
308
309    # Warn about corrupt PPM files
310    for f in glob.glob(os.path.join(test.debugdir, "*.ppm")):
311        if not ppm_utils.image_verify_ppm_file(f):
312            logging.warning("Found corrupt PPM file: %s", f)
313
314    # Should we convert PPM files to PNG format?
315    if params.get("convert_ppm_files_to_png") == "yes":
316        logging.debug("Param 'convert_ppm_files_to_png' specified, converting "
317                      "PPM files to PNG format")
318        try:
319            for f in glob.glob(os.path.join(test.debugdir, "*.ppm")):
320                if ppm_utils.image_verify_ppm_file(f):
321                    new_path = f.replace(".ppm", ".png")
322                    image = PIL.Image.open(f)
323                    image.save(new_path, format='PNG')
324        except NameError:
325            pass
326
327    # Should we keep the PPM files?
328    if params.get("keep_ppm_files") != "yes":
329        logging.debug("Param 'keep_ppm_files' not specified, removing all PPM "
330                      "files from debug dir")
331        for f in glob.glob(os.path.join(test.debugdir, '*.ppm')):
332            os.unlink(f)
333
334    # Should we keep the screendump dirs?
335    if params.get("keep_screendumps") != "yes":
336        logging.debug("Param 'keep_screendumps' not specified, removing "
337                      "screendump dirs")
338        for d in glob.glob(os.path.join(test.debugdir, "screendumps_*")):
339            if os.path.isdir(d) and not os.path.islink(d):
340                shutil.rmtree(d, ignore_errors=True)
341
342    # Kill all unresponsive VMs
343    if params.get("kill_unresponsive_vms") == "yes":
344        logging.debug("Param 'kill_unresponsive_vms' specified, killing all "
345                      "VMs that fail to respond to a remote login request")
346        for vm in env.get_all_vms():
347            if vm.is_alive():
348                try:
349                    session = vm.login()
350                    session.close()
351                except (virt_utils.LoginError, virt_vm.VMError), e:
352                    logging.warning(e)
353                    vm.destroy(gracefully=False)
354
355    # Kill all aexpect tail threads
356    aexpect.kill_tail_threads()
357
358    # Terminate tcpdump if no VMs are alive
359    living_vms = [vm for vm in env.get_all_vms() if vm.is_alive()]
360    if not living_vms and "tcpdump" in env:
361        env["tcpdump"].close()
362        del env["tcpdump"]
363
364    if params.get("setup_hugepages") == "yes":
365        h = virt_test_setup.HugePageConfig(params)
366        h.cleanup()
367
368    # Execute any post_commands
369    if params.get("post_command"):
370        process_command(test, params, env, params.get("post_command"),
371                        int(params.get("post_command_timeout", "600")),
372                        params.get("post_command_noncritical") == "yes")
373
374
375def postprocess_on_error(test, params, env):
376    """
377    Perform postprocessing operations required only if the test failed.
378
379    @param test: An Autotest test object.
380    @param params: A dict containing all VM and image parameters.
381    @param env: The environment (a dict-like object).
382    """
383    params.update(params.object_params("on_error"))
384
385
386def _update_address_cache(address_cache, line):
387    if re.search("Your.IP", line, re.IGNORECASE):
388        matches = re.findall(r"\d*\.\d*\.\d*\.\d*", line)
389        if matches:
390            address_cache["last_seen"] = matches[0]
391    if re.search("Client.Ethernet.Address", line, re.IGNORECASE):
392        matches = re.findall(r"\w*:\w*:\w*:\w*:\w*:\w*", line)
393        if matches and address_cache.get("last_seen"):
394            mac_address = matches[0].lower()
395            if time.time() - address_cache.get("time_%s" % mac_address, 0) > 5:
396                logging.debug("(address cache) Adding cache entry: %s ---> %s",
397                              mac_address, address_cache.get("last_seen"))
398            address_cache[mac_address] = address_cache.get("last_seen")
399            address_cache["time_%s" % mac_address] = time.time()
400            del address_cache["last_seen"]
401
402
403def _take_screendumps(test, params, env):
404    global _screendump_thread_termination_event
405    temp_dir = test.debugdir
406    if params.get("screendump_temp_dir"):
407        temp_dir = virt_utils.get_path(test.bindir,
408                                      params.get("screendump_temp_dir"))
409        try:
410            os.makedirs(temp_dir)
411        except OSError:
412            pass
413    temp_filename = os.path.join(temp_dir, "scrdump-%s.ppm" %
414                                 virt_utils.generate_random_string(6))
415    delay = float(params.get("screendump_delay", 5))
416    quality = int(params.get("screendump_quality", 30))
417
418    cache = {}
419
420    while True:
421        for vm in env.get_all_vms():
422            if not vm.is_alive():
423                continue
424            try:
425                vm.monitor.screendump(filename=temp_filename, debug=False)
426            except kvm_monitor.MonitorError, e:
427                logging.warning(e)
428                continue
429            except AttributeError, e:
430                continue
431            if not os.path.exists(temp_filename):
432                logging.warning("VM '%s' failed to produce a screendump", vm.name)
433                continue
434            if not ppm_utils.image_verify_ppm_file(temp_filename):
435                logging.warning("VM '%s' produced an invalid screendump", vm.name)
436                os.unlink(temp_filename)
437                continue
438            screendump_dir = os.path.join(test.debugdir,
439                                          "screendumps_%s" % vm.name)
440            try:
441                os.makedirs(screendump_dir)
442            except OSError:
443                pass
444            screendump_filename = os.path.join(screendump_dir,
445                    "%s_%s.jpg" % (vm.name,
446                                   time.strftime("%Y-%m-%d_%H-%M-%S")))
447            hash = utils.hash_file(temp_filename)
448            if hash in cache:
449                try:
450                    os.link(cache[hash], screendump_filename)
451                except OSError:
452                    pass
453            else:
454                try:
455                    image = PIL.Image.open(temp_filename)
456                    image.save(screendump_filename, format="JPEG", quality=quality)
457                    cache[hash] = screendump_filename
458                except NameError:
459                    pass
460            os.unlink(temp_filename)
461        if _screendump_thread_termination_event.isSet():
462            _screendump_thread_termination_event = None
463            break
464        _screendump_thread_termination_event.wait(delay)
465