htest/vm/controller.py

   1 #!/usr/bin/env python3
   2
   3 #
   4 # Copyright (c) 2018 Vojtech Horky
   5 # All rights reserved.
   6 #
   7 # Redistribution and use in source and binary forms, with or without
   8 # modification, are permitted provided that the following conditions
   9 # are met:
  10 #
  11 # - Redistributions of source code must retain the above copyright
  12 #   notice, this list of conditions and the following disclaimer.
  13 # - Redistributions in binary form must reproduce the above copyright
  14 #   notice, this list of conditions and the following disclaimer in the
  15 #   documentation and/or other materials provided with the distribution.
  16 # - The name of the author may not be used to endorse or promote products
  17 #   derived from this software without specific prior written permission.
  18 #
  19 # THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
  20 # IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
  21 # OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
  22 # IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
  23 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
  24 # NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  25 # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  26 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  27 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
  28 # THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  29 #
  30
  31 import os
  32 import subprocess
  33
  34 class VMManager:
  35     """
  36     Keeps track of running virtual machines.
  37     """
  38
  39     def __init__(self, controller, architecture, boot_image, memory_amount):
  40         self.controller_class = controller
  41         self.architecture = architecture
  42         self.boot_image = boot_image
  43         self.memory_amount = memory_amount
  44         self.instances = {}
  45         self.last = None
  46
  47     def create(self, name):
  48         if name in self.instances:
  49             raise Exception("Duplicate machine name {}.".format(name))
  50         self.instances[name] = self.controller_class(self.architecture, name, self.boot_image)
  51         self.instances[name].memory = self.memory_amount
  52         self.last = name
  53         return self.instances[name]
  54
  55     def get(self, name=None):
  56         if name is None:
  57             name = self.last
  58         if name is None:
  59             return None
  60         if name in self.instances:
  61             self.last = name
  62             return self.instances[name]
  63         else:
  64             return None
  65
  66     def terminate(self, vterm_dump_filename, last_screenshot_filename):
  67         for i in self.instances:
  68             self.instances[i].terminate()
  69         if vterm_dump_filename is not None:
  70             with open(vterm_dump_filename, 'w') as f:
  71                 for i in self.instances:
  72                     lines = '\n'.join(self.instances[i].full_vterm)
  73                     print(lines, file=f)
  74         if last_screenshot_filename is not None:
  75             for i in self.instances:
  76                 filename = self.instances[i].screenshot_filename
  77                 if filename is not None:
  78                     proc = subprocess.Popen(['convert', filename, last_screenshot_filename])
  79                     proc.wait()
  80                     if proc.returncode != 0:
  81                         raise Exception("Saving screenshot failed.")
  82
  83 class VMController:
  84     """
  85     Base class for controllers of specific virtual machine emulators.
  86     """
  87
  88     def __init__(self, provider):
  89         self.provider_name = provider
  90         # Patched by VMManager
  91         self.name = 'XXX'
  92         self.screenshot_filename = None
  93         # All lines seen in the terminal
  94         # (do not reset unless you know what you are doing).
  95         self.full_vterm = []
  96         # Used to keep track of new-lines
  97         self.vterm = []
  98         # Amount of memory (MB) (patched by VMM manager)
  99         self.memory = 0
 100         pass
 101
 102     def is_supported(self, arch):
 103         """
 104         Tells whether this controller supports given architecture.
 105         """
 106         return False
 107
 108     def boot(self, **kwargs):
 109         """
 110         Bring the machine up.
 111         """
 112         pass
 113
 114     def terminate(self):
 115         """
 116         Shutdown the VM.
 117         """
 118         pass
 119
 120     def type(self, what):
 121         """
 122         Type given text into vterm.
 123         """
 124         print("type('{}') @ {}".format(what, self.provider_name))
 125         pass
 126
 127     def same_vterm_tail(self, lines):
 128         lines_count = len(lines)
 129         for i in range(-1, -lines_count - 1, -1):
 130             if i != -1:
 131                 if lines[i] != self.full_vterm[i]:
 132                     return False
 133             else:
 134                 a = lines[-1].replace("_", " ").strip()
 135                 b = self.full_vterm[-1].replace("_", " ").strip()
 136                 if not a.startswith(b):
 137                     return False
 138         return True
 139
 140     def capture_vterm(self):
 141         """
 142         Capture contents of current terminal window and updates self.vterm
 143         """
 144
 145         # Read everything from the terminal and get rid of completely empty
 146         # lines (for first commands when the screen is empty).
 147         lines = self.capture_vterm_impl()
 148         lines = [l.strip() for l in lines]
 149         while (len(lines) > 0) and (lines[-1].strip() == ""):
 150             lines = lines[0:-1]
 151         if (len(lines) == 0):
 152             return
 153
 154         # When this is the very first screen, we simply copy it.
 155         if len(self.full_vterm) == 0:
 156             for l in lines:
 157                 self.full_vterm.append(l)
 158                 self.vterm.append(l)
 159         else:
 160             # Otherwise, we find whether there is some overlap, i.e. whether
 161             # we are capturing a rolling screen.
 162             lines_count = len(lines)
 163             same_lines = 0
 164             for i in range(lines_count, 0, -1):
 165                 if self.same_vterm_tail(lines[0:i]):
 166                     same_lines = i
 167                     break
 168             # If there is no overlap, we might have missed some lines.
 169             if same_lines == 0:
 170                 self.full_vterm.append("!!!!!! WARNING: probably missed some lines here !!!!!")
 171             else:
 172                 # Otherwise, update the last line (last capture might have
 173                 # missed some characters).
 174                 if len(self.full_vterm) > 0:
 175                     self.full_vterm = self.full_vterm[0:-1]
 176                 if len(self.vterm) > 0:
 177                     self.vterm = self.vterm[0:-1]
 178                 same_lines = same_lines - 1
 179             # Add the new lines.
 180             for i in range(same_lines, lines_count):
 181                 self.full_vterm.append(lines[i])
 182                 self.vterm.append(lines[i])
 183
 184     def capture_vterm_impl(self):
 185         """
 186         Do not call but reimplement in subclass.
 187         """
 188         return []
 189
 190     def get_temp(self, id):
 191         """
 192         Get temporary file name.
 193         """
 194         os.makedirs('tmp-vm-python', exist_ok=True)
 195         return 'tmp-vm-python/tmp-' + self.name + '-' + id
 196