xref: /openbmc/qemu/tests/vm/basevm.py (revision f3635813) 
1#!/usr/bin/env python
2#
3# VM testing base class
4#
5# Copyright 2017-2019 Red Hat Inc.
6#
7# Authors:
8#  Fam Zheng <famz@redhat.com>
9#  Gerd Hoffmann <kraxel@redhat.com>
10#
11# This code is licensed under the GPL version 2 or later.  See
12# the COPYING file in the top-level directory.
13#
14
15from __future__ import print_function
16import os
17import re
18import sys
19import socket
20import logging
21import time
22import datetime
23sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'python'))
24from qemu import kvm_available
25from qemu.machine import QEMUMachine
26import subprocess
27import hashlib
28import optparse
29import atexit
30import tempfile
31import shutil
32import multiprocessing
33import traceback
34
35SSH_KEY = open(os.path.join(os.path.dirname(__file__),
36               "..", "keys", "id_rsa")).read()
37SSH_PUB_KEY = open(os.path.join(os.path.dirname(__file__),
38                   "..", "keys", "id_rsa.pub")).read()
39
40class BaseVM(object):
41    GUEST_USER = "qemu"
42    GUEST_PASS = "qemupass"
43    ROOT_PASS = "qemupass"
44
45    envvars = [
46        "https_proxy",
47        "http_proxy",
48        "ftp_proxy",
49        "no_proxy",
50    ]
51
52    # The script to run in the guest that builds QEMU
53    BUILD_SCRIPT = ""
54    # The guest name, to be overridden by subclasses
55    name = "#base"
56    # The guest architecture, to be overridden by subclasses
57    arch = "#arch"
58    # command to halt the guest, can be overridden by subclasses
59    poweroff = "poweroff"
60    # enable IPv6 networking
61    ipv6 = True
62    def __init__(self, debug=False, vcpus=None):
63        self._guest = None
64        self._tmpdir = os.path.realpath(tempfile.mkdtemp(prefix="vm-test-",
65                                                         suffix=".tmp",
66                                                         dir="."))
67        atexit.register(shutil.rmtree, self._tmpdir)
68
69        self._ssh_key_file = os.path.join(self._tmpdir, "id_rsa")
70        open(self._ssh_key_file, "w").write(SSH_KEY)
71        subprocess.check_call(["chmod", "600", self._ssh_key_file])
72
73        self._ssh_pub_key_file = os.path.join(self._tmpdir, "id_rsa.pub")
74        open(self._ssh_pub_key_file, "w").write(SSH_PUB_KEY)
75
76        self.debug = debug
77        self._stderr = sys.stderr
78        self._devnull = open(os.devnull, "w")
79        if self.debug:
80            self._stdout = sys.stdout
81        else:
82            self._stdout = self._devnull
83        self._args = [ \
84            "-nodefaults", "-m", "4G",
85            "-cpu", "max",
86            "-netdev", "user,id=vnet,hostfwd=:127.0.0.1:0-:22" +
87                       (",ipv6=no" if not self.ipv6 else ""),
88            "-device", "virtio-net-pci,netdev=vnet",
89            "-vnc", "127.0.0.1:0,to=20"]
90        if vcpus and vcpus > 1:
91            self._args += ["-smp", "%d" % vcpus]
92        if kvm_available(self.arch):
93            self._args += ["-enable-kvm"]
94        else:
95            logging.info("KVM not available, not using -enable-kvm")
96        self._data_args = []
97
98    def _download_with_cache(self, url, sha256sum=None, sha512sum=None):
99        def check_sha256sum(fname):
100            if not sha256sum:
101                return True
102            checksum = subprocess.check_output(["sha256sum", fname]).split()[0]
103            return sha256sum == checksum.decode("utf-8")
104
105        def check_sha512sum(fname):
106            if not sha512sum:
107                return True
108            checksum = subprocess.check_output(["sha512sum", fname]).split()[0]
109            return sha512sum == checksum.decode("utf-8")
110
111        cache_dir = os.path.expanduser("~/.cache/qemu-vm/download")
112        if not os.path.exists(cache_dir):
113            os.makedirs(cache_dir)
114        fname = os.path.join(cache_dir,
115                             hashlib.sha1(url.encode("utf-8")).hexdigest())
116        if os.path.exists(fname) and check_sha256sum(fname) and check_sha512sum(fname):
117            return fname
118        logging.debug("Downloading %s to %s...", url, fname)
119        subprocess.check_call(["wget", "-c", url, "-O", fname + ".download"],
120                              stdout=self._stdout, stderr=self._stderr)
121        os.rename(fname + ".download", fname)
122        return fname
123
124    def _ssh_do(self, user, cmd, check):
125        ssh_cmd = ["ssh", "-q", "-t",
126                   "-o", "StrictHostKeyChecking=no",
127                   "-o", "UserKnownHostsFile=" + os.devnull,
128                   "-o", "ConnectTimeout=1",
129                   "-p", self.ssh_port, "-i", self._ssh_key_file]
130        for var in self.envvars:
131            ssh_cmd += ['-o', "SendEnv=%s" % var ]
132        assert not isinstance(cmd, str)
133        ssh_cmd += ["%s@127.0.0.1" % user] + list(cmd)
134        logging.debug("ssh_cmd: %s", " ".join(ssh_cmd))
135        r = subprocess.call(ssh_cmd)
136        if check and r != 0:
137            raise Exception("SSH command failed: %s" % cmd)
138        return r
139
140    def ssh(self, *cmd):
141        return self._ssh_do(self.GUEST_USER, cmd, False)
142
143    def ssh_root(self, *cmd):
144        return self._ssh_do("root", cmd, False)
145
146    def ssh_check(self, *cmd):
147        self._ssh_do(self.GUEST_USER, cmd, True)
148
149    def ssh_root_check(self, *cmd):
150        self._ssh_do("root", cmd, True)
151
152    def build_image(self, img):
153        raise NotImplementedError
154
155    def add_source_dir(self, src_dir):
156        name = "data-" + hashlib.sha1(src_dir.encode("utf-8")).hexdigest()[:5]
157        tarfile = os.path.join(self._tmpdir, name + ".tar")
158        logging.debug("Creating archive %s for src_dir dir: %s", tarfile, src_dir)
159        subprocess.check_call(["./scripts/archive-source.sh", tarfile],
160                              cwd=src_dir, stdin=self._devnull,
161                              stdout=self._stdout, stderr=self._stderr)
162        self._data_args += ["-drive",
163                            "file=%s,if=none,id=%s,cache=writeback,format=raw" % \
164                                    (tarfile, name),
165                            "-device",
166                            "virtio-blk,drive=%s,serial=%s,bootindex=1" % (name, name)]
167
168    def boot(self, img, extra_args=[]):
169        args = self._args + [
170            "-device", "VGA",
171            "-drive", "file=%s,if=none,id=drive0,cache=writeback" % img,
172            "-device", "virtio-blk,drive=drive0,bootindex=0"]
173        args += self._data_args + extra_args
174        logging.debug("QEMU args: %s", " ".join(args))
175        qemu_bin = os.environ.get("QEMU", "qemu-system-" + self.arch)
176        guest = QEMUMachine(binary=qemu_bin, args=args)
177        guest.set_machine('pc')
178        guest.set_console()
179        try:
180            guest.launch()
181        except:
182            logging.error("Failed to launch QEMU, command line:")
183            logging.error(" ".join([qemu_bin] + args))
184            logging.error("Log:")
185            logging.error(guest.get_log())
186            logging.error("QEMU version >= 2.10 is required")
187            raise
188        atexit.register(self.shutdown)
189        self._guest = guest
190        usernet_info = guest.qmp("human-monitor-command",
191                                 command_line="info usernet")
192        self.ssh_port = None
193        for l in usernet_info["return"].splitlines():
194            fields = l.split()
195            if "TCP[HOST_FORWARD]" in fields and "22" in fields:
196                self.ssh_port = l.split()[3]
197        if not self.ssh_port:
198            raise Exception("Cannot find ssh port from 'info usernet':\n%s" % \
199                            usernet_info)
200
201    def console_init(self, timeout = 120):
202        vm = self._guest
203        vm.console_socket.settimeout(timeout)
204
205    def console_log(self, text):
206        for line in re.split("[\r\n]", text):
207            # filter out terminal escape sequences
208            line = re.sub("\x1b\[[0-9;?]*[a-zA-Z]", "", line)
209            line = re.sub("\x1b\([0-9;?]*[a-zA-Z]", "", line)
210            # replace unprintable chars
211            line = re.sub("\x1b", "<esc>", line)
212            line = re.sub("[\x00-\x1f]", ".", line)
213            line = re.sub("[\x80-\xff]", ".", line)
214            if line == "":
215                continue
216            # log console line
217            sys.stderr.write("con recv: %s\n" % line)
218
219    def console_wait(self, expect, expectalt = None):
220        vm = self._guest
221        output = ""
222        while True:
223            try:
224                chars = vm.console_socket.recv(1)
225            except socket.timeout:
226                sys.stderr.write("console: *** read timeout ***\n")
227                sys.stderr.write("console: waiting for: '%s'\n" % expect)
228                if not expectalt is None:
229                    sys.stderr.write("console: waiting for: '%s' (alt)\n" % expectalt)
230                sys.stderr.write("console: line buffer:\n")
231                sys.stderr.write("\n")
232                self.console_log(output.rstrip())
233                sys.stderr.write("\n")
234                raise
235            output += chars.decode("latin1")
236            if expect in output:
237                break
238            if not expectalt is None and expectalt in output:
239                break
240            if "\r" in output or "\n" in output:
241                lines = re.split("[\r\n]", output)
242                output = lines.pop()
243                if self.debug:
244                    self.console_log("\n".join(lines))
245        if self.debug:
246            self.console_log(output)
247        if not expectalt is None and expectalt in output:
248            return False
249        return True
250
251    def console_consume(self):
252        vm = self._guest
253        output = ""
254        vm.console_socket.setblocking(0)
255        while True:
256            try:
257                chars = vm.console_socket.recv(1)
258            except:
259                break
260            output += chars.decode("latin1")
261            if "\r" in output or "\n" in output:
262                lines = re.split("[\r\n]", output)
263                output = lines.pop()
264                if self.debug:
265                    self.console_log("\n".join(lines))
266        if self.debug:
267            self.console_log(output)
268        vm.console_socket.setblocking(1)
269
270    def console_send(self, command):
271        vm = self._guest
272        if self.debug:
273            logline = re.sub("\n", "<enter>", command)
274            logline = re.sub("[\x00-\x1f]", ".", logline)
275            sys.stderr.write("con send: %s\n" % logline)
276        for char in list(command):
277            vm.console_socket.send(char.encode("utf-8"))
278            time.sleep(0.01)
279
280    def console_wait_send(self, wait, command):
281        self.console_wait(wait)
282        self.console_send(command)
283
284    def console_ssh_init(self, prompt, user, pw):
285        sshkey_cmd = "echo '%s' > .ssh/authorized_keys\n" % SSH_PUB_KEY.rstrip()
286        self.console_wait_send("login:",    "%s\n" % user)
287        self.console_wait_send("Password:", "%s\n" % pw)
288        self.console_wait_send(prompt,      "mkdir .ssh\n")
289        self.console_wait_send(prompt,      sshkey_cmd)
290        self.console_wait_send(prompt,      "chmod 755 .ssh\n")
291        self.console_wait_send(prompt,      "chmod 644 .ssh/authorized_keys\n")
292
293    def console_sshd_config(self, prompt):
294        self.console_wait(prompt)
295        self.console_send("echo 'PermitRootLogin yes' >> /etc/ssh/sshd_config\n")
296        for var in self.envvars:
297            self.console_wait(prompt)
298            self.console_send("echo 'AcceptEnv %s' >> /etc/ssh/sshd_config\n" % var)
299
300    def print_step(self, text):
301        sys.stderr.write("### %s ...\n" % text)
302
303    def wait_ssh(self, seconds=300):
304        starttime = datetime.datetime.now()
305        endtime = starttime + datetime.timedelta(seconds=seconds)
306        guest_up = False
307        while datetime.datetime.now() < endtime:
308            if self.ssh("exit 0") == 0:
309                guest_up = True
310                break
311            seconds = (endtime - datetime.datetime.now()).total_seconds()
312            logging.debug("%ds before timeout", seconds)
313            time.sleep(1)
314        if not guest_up:
315            raise Exception("Timeout while waiting for guest ssh")
316
317    def shutdown(self):
318        self._guest.shutdown()
319
320    def wait(self):
321        self._guest.wait()
322
323    def graceful_shutdown(self):
324        self.ssh_root(self.poweroff)
325        self._guest.wait()
326
327    def qmp(self, *args, **kwargs):
328        return self._guest.qmp(*args, **kwargs)
329
330def parse_args(vmcls):
331
332    def get_default_jobs():
333        if kvm_available(vmcls.arch):
334            return multiprocessing.cpu_count() // 2
335        else:
336            return 1
337
338    parser = optparse.OptionParser(
339        description="VM test utility.  Exit codes: "
340                    "0 = success, "
341                    "1 = command line error, "
342                    "2 = environment initialization failed, "
343                    "3 = test command failed")
344    parser.add_option("--debug", "-D", action="store_true",
345                      help="enable debug output")
346    parser.add_option("--image", "-i", default="%s.img" % vmcls.name,
347                      help="image file name")
348    parser.add_option("--force", "-f", action="store_true",
349                      help="force build image even if image exists")
350    parser.add_option("--jobs", type=int, default=get_default_jobs(),
351                      help="number of virtual CPUs")
352    parser.add_option("--verbose", "-V", action="store_true",
353                      help="Pass V=1 to builds within the guest")
354    parser.add_option("--build-image", "-b", action="store_true",
355                      help="build image")
356    parser.add_option("--build-qemu",
357                      help="build QEMU from source in guest")
358    parser.add_option("--build-target",
359                      help="QEMU build target", default="check")
360    parser.add_option("--interactive", "-I", action="store_true",
361                      help="Interactively run command")
362    parser.add_option("--snapshot", "-s", action="store_true",
363                      help="run tests with a snapshot")
364    parser.disable_interspersed_args()
365    return parser.parse_args()
366
367def main(vmcls):
368    try:
369        args, argv = parse_args(vmcls)
370        if not argv and not args.build_qemu and not args.build_image:
371            print("Nothing to do?")
372            return 1
373        logging.basicConfig(level=(logging.DEBUG if args.debug
374                                   else logging.WARN))
375        vm = vmcls(debug=args.debug, vcpus=args.jobs)
376        if args.build_image:
377            if os.path.exists(args.image) and not args.force:
378                sys.stderr.writelines(["Image file exists: %s\n" % args.image,
379                                      "Use --force option to overwrite\n"])
380                return 1
381            return vm.build_image(args.image)
382        if args.build_qemu:
383            vm.add_source_dir(args.build_qemu)
384            cmd = [vm.BUILD_SCRIPT.format(
385                   configure_opts = " ".join(argv),
386                   jobs=int(args.jobs),
387                   target=args.build_target,
388                   verbose = "V=1" if args.verbose else "")]
389        else:
390            cmd = argv
391        img = args.image
392        if args.snapshot:
393            img += ",snapshot=on"
394        vm.boot(img)
395        vm.wait_ssh()
396    except Exception as e:
397        if isinstance(e, SystemExit) and e.code == 0:
398            return 0
399        sys.stderr.write("Failed to prepare guest environment\n")
400        traceback.print_exc()
401        return 2
402
403    exitcode = 0
404    if vm.ssh(*cmd) != 0:
405        exitcode = 3
406    if args.interactive:
407        vm.ssh()
408
409    if not args.snapshot:
410        vm.graceful_shutdown()
411
412    return exitcode
413