Add support for shrinking windows VMs
[snf-image-creator] / image_creator / os_type / windows.py
index c78c947..f8e30f6 100644 (file)
@@ -46,6 +46,7 @@ import time
 import random
 import string
 import subprocess
+import struct
 
 kvm = get_command('kvm')
 
@@ -54,6 +55,22 @@ BOOT_TIMEOUT = 300
 
 class Windows(OSBase):
     """OS class for Windows"""
+    def __init__(self, image, **kargs):
+        super(Windows, self).__init__(image, **kargs)
+
+        device = self.g.part_to_dev(self.root)
+
+        self.last_part_num = self.g.part_list(device)[-1]['part_num']
+        self.last_drive = None
+        self.system_drive = None
+
+        for drive, partition in self.g.inspect_get_drive_mappings(self.root):
+            if partition == "%s%d" % (device, self.last_part_num):
+                self.last_drive = drive
+            if partition == self.root:
+                self.system_drive = drive
+
+        assert self.system_drive
 
     def needed_sysprep_params(self):
         """Returns a list of needed sysprep parameters. Each element in the
@@ -89,6 +106,12 @@ class Windows(OSBase):
 
         self._guest_exec('netsh firewall set icmpsetting 8')
 
+    @sysprep('Disabling hibernation support')
+    def disable_hibernation(self):
+        """Disable hibernation support and remove the hibernation file"""
+
+        self._guest_exec(r'powercfg.exe /hibernate off')
+
     @sysprep('Setting the system clock to UTC')
     def utc(self):
         """Set the hardware clock to UTC"""
@@ -97,6 +120,13 @@ class Windows(OSBase):
         self._guest_exec(
             r'REG ADD %s /v RealTimeIsUniversal /t REG_DWORD /d 1 /f' % path)
 
+    @sysprep('Clearing the event logs')
+    def clear_logs(self):
+        """Clear all the event logs"""
+
+        self._guest_exec(
+            r"cmd /q /c for /f %l in ('wevtutil el') do wevtutil cl %l")
+
     @sysprep('Executing sysprep on the image (may take more that 10 minutes)')
     def microsoft_sysprep(self):
         """Run the Microsoft System Preparation Tool. This will remove
@@ -108,6 +138,66 @@ class Windows(OSBase):
                          r'/quiet /generalize /oobe /shutdown')
         self.syspreped = True
 
+    @sysprep('Shrinking the last filesystem')
+    def shrink(self):
+        """Shrink the last filesystem. Make sure the filesystem is defragged"""
+
+        # Query for the maximum number of reclaimable bytes
+        cmd = (
+            r'cmd /Q /C "SET SCRIPT=%TEMP%\QUERYMAX_%RANDOM%.TXT & ' +
+            r'ECHO SELECT DISK 0 > %SCRIPT% & ' +
+            'ECHO SELECT PARTITION %d >> %%SCRIPT%% & ' % self.last_part_num +
+            r'ECHO SHRINK QUERYMAX >> %SCRIPT% & ' +
+            r'ECHO EXIT >> %SCRIPT% & ' +
+            r'DISKPART /S %SCRIPT% & ' +
+            r'IF ERRORLEVEL 1 EXIT /B 1 & ' +
+            r'DEL /Q %SCRIPT%"')
+
+        stdout, stderr, rc = self._guest_exec(cmd)
+
+        querymax = None
+        for line in stdout.splitlines():
+            # diskpart will return something like this:
+            #
+            #   The maximum number of reclaimable bytes is: xxxx MB
+            #
+            if line.find('reclaimable') >= 0:
+                querymax = line.split(':')[1].split()[0].strip()
+                assert querymax.isdigit(), \
+                    "Number of reclaimable bytes not a number"
+
+        if querymax is None:
+            FatalError("Error in shrinking! "
+                       "Couldn't find the max number of reclaimable bytes!")
+
+        querymax = int(querymax)
+        # From ntfsresize:
+        # Practically the smallest shrunken size generally is at around
+        # "used space" + (20-200 MB). Please also take into account that
+        # Windows might need about 50-100 MB free space left to boot safely.
+        # I'll give 100MB extra space just to be sure
+        querymax -= 100
+
+        if querymax < 0:
+            self.out.warn("Not enought available space to shrink the image!")
+            return
+
+        cmd = (
+            r'cmd /Q /C "SET SCRIPT=%TEMP%\QUERYMAX_%RANDOM%.TXT & ' +
+            r'ECHO SELECT DISK 0 > %SCRIPT% & ' +
+            'ECHO SELECT PARTITION %d >> %%SCRIPT%% & ' % self.last_part_num +
+            'ECHO SHRINK DESIRED=%d >> %%SCRIPT%% & ' % querymax +
+            r'ECHO EXIT >> %SCRIPT% & ' +
+            r'DISKPART /S %SCRIPT% & ' +
+            r'IF ERRORLEVEL 1 EXIT /B 1 & ' +
+            r'DEL /Q %SCRIPT%"')
+
+        stdout, stderr, rc = self._guest_exec(cmd)
+
+        for line in stdout.splitlines():
+            if line.find('shrunk') >= 0:
+                self.out.output(line)
+
     def do_sysprep(self):
         """Prepare system for image creation."""
 
@@ -123,6 +213,15 @@ class Windows(OSBase):
         try:
             disabled_uac = self._update_uac_remote_setting(1)
             token = self._enable_os_monitor()
+
+            # disable the firewalls
+            firewall_states = self._update_firewalls(0, 0, 0)
+
+            # Delete the pagefile. It will be recreated when the system boots
+            systemroot = self.g.inspect_get_windows_systemroot(self.root)
+            pagefile = "%s/pagefile.sys" % systemroot
+            self.g.rm_rf(self.g.case_sensitive_path(pagefile))
+
         finally:
             self.umount()
 
@@ -152,6 +251,8 @@ class Windows(OSBase):
             else:
                 self.out.success('done')
 
+            time.sleep(5)  # Just to be sure everything is up
+
             self.out.output("Disabling automatic logon ...", False)
             self._disable_autologon()
             self.out.success('done')
@@ -160,12 +261,21 @@ class Windows(OSBase):
 
             tasks = self.list_syspreps()
             enabled = filter(lambda x: x.enabled, tasks)
-
             size = len(enabled)
 
+            # Make sure shrink runs in the end, before ms sysprep
+            enabled = filter(lambda x: self.sysprep_info(x).name != 'shrink',
+                             enabled)
+
+            shrink_enabled = False
+            if len(enabled) != size:
+                enabled.append(self.shrink)
+                shrink_enabled = True
+
             # Make sure the ms sysprep is the last task to run if it is enabled
             enabled = filter(
-                lambda x: x.im_func.func_name != 'microsoft_sysprep', enabled)
+                lambda x: self.sysprep_info(x).name != 'microsoft-sysprep',
+                enabled)
 
             ms_sysprep_enabled = False
             if len(enabled) != size:
@@ -179,12 +289,14 @@ class Windows(OSBase):
                 task()
                 setattr(task.im_func, 'executed', True)
 
-            self.out.output("Shutting down windows VM ...", False)
+            self.out.output("Sending shut down command ...", False)
             if not ms_sysprep_enabled:
                 self._shutdown()
             self.out.success("done")
 
+            self.out.output("Waiting for windows to shut down ...", False)
             vm.wait()
+            self.out.success("done")
         finally:
             if monitor is not None:
                 os.unlink(monitor)
@@ -197,8 +309,14 @@ class Windows(OSBase):
             self.g.launch()
             self.out.success('done')
 
-        if disabled_uac:
-            self._update_uac_remote_setting(0)
+            self.mount(readonly=False)
+            try:
+                if disabled_uac:
+                    self._update_uac_remote_setting(0)
+
+                self._update_firewalls(*firewall_states)
+            finally:
+                self.umount()
 
     def _create_vm(self, monitor):
         """Create a VM with the image attached as the disk
@@ -218,12 +336,12 @@ class Windows(OSBase):
         vnc_port = random.randint(11000, 14999)
         display = vnc_port - 5900
 
-        vm = kvm('-smp', '1', '-m', '1024', '-drive',
-                 'file=%s,format=raw,cache=none,if=virtio' % self.image.device,
-                 '-netdev', 'type=user,hostfwd=tcp::445-:445,id=netdev0',
-                 '-device', 'virtio-net-pci,mac=%s,netdev=netdev0' %
-                 random_mac(), '-vnc', ':%d' % display, '-serial',
-                 'file:%s' % monitor, _bg=True)
+        vm = kvm(
+            '-smp', '1', '-m', '1024', '-drive',
+            'file=%s,format=raw,cache=unsafe,if=virtio' % self.image.device,
+            '-netdev', 'type=user,hostfwd=tcp::445-:445,id=netdev0',
+            '-device', 'virtio-net-pci,mac=%s,netdev=netdev0' % random_mac(),
+            '-vnc', ':%d' % display, '-serial', 'file:%s' % monitor, _bg=True)
 
         return vm, display
 
@@ -335,6 +453,14 @@ class Windows(OSBase):
             h.node_set_value(runonce,
                              {'key': "BootMonitor", 't': 1, 'value': value})
 
+            value = (
+                r'REG ADD HKLM\SOFTWARE\Microsoft\Windows\CurrentVersion'
+                r'\policies\system /v LocalAccountTokenFilterPolicy'
+                r' /t REG_DWORD /d 1 /f').encode('utf-16le')
+
+            h.node_set_value(runonce,
+                             {'key': "UpdateRegistry", 't': 1, 'value': value})
+
             h.commit(None)
 
             self.g.upload(software, path)
@@ -343,6 +469,65 @@ class Windows(OSBase):
 
         return token
 
+    def _update_firewalls(self, domain, public, standard):
+        """Enables or disables the firewall for the Domain, the Public and the
+        Standard profile. Returns a triplete with the old values.
+
+        1 will enable a firewall and 0 will disable it
+        """
+
+        if domain not in (0, 1):
+            raise ValueError("Valid values for domain parameter are 0 and 1")
+
+        if public not in (0, 1):
+            raise ValueError("Valid values for public parameter are 0 and 1")
+
+        if standard not in (0, 1):
+            raise ValueError("Valid values for standard parameter are 0 and 1")
+
+        path = self._registry_file_path("SYSTEM")
+        systemfd, system = tempfile.mkstemp()
+        try:
+            os.close(systemfd)
+            self.g.download(path, system)
+
+            h = hivex.Hivex(system, write=True)
+
+            select = h.node_get_child(h.root(), 'Select')
+            current_value = h.node_get_value(select, 'Current')
+
+            # expecting a little endian dword
+            assert h.value_type(current_value)[1] == 4
+            current = "%03d" % h.value_dword(current_value)
+
+            firewall_policy = h.root()
+            for child in ('ControlSet%s' % current, 'services', 'SharedAccess',
+                          'Parameters', 'FirewallPolicy'):
+                firewall_policy = h.node_get_child(firewall_policy, child)
+
+            old_values = []
+            new_values = [domain, public, standard]
+            for profile in ('Domain', 'Public', 'Standard'):
+                node = h.node_get_child(firewall_policy, '%sProfile' % profile)
+
+                old_value = h.node_get_value(node, 'EnableFirewall')
+
+                # expecting a little endian dword
+                assert h.value_type(old_value)[1] == 4
+                old_values.append(h.value_dword(old_value))
+
+                h.node_set_value(
+                    node, {'key': 'EnableFirewall', 't': 4L,
+                           'value': struct.pack("<I", new_values.pop(0))})
+
+            h.commit(None)
+            self.g.upload(system, path)
+
+        finally:
+            os.unlink(system)
+
+        return old_values
+
     def _update_uac_remote_setting(self, value):
         """Updates the registry key value:
         [HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows\CurrentVersion\Policies
@@ -386,9 +571,8 @@ class Windows(OSBase):
             elif value == 0:
                 return False
 
-            new_value = {
-                'key': "LocalAccountTokenFilterPolicy", 't': 4L,
-                'value': '%s\x00\x00\x00' % '\x00' if value == 0 else '\x01'}
+            new_value = {'key': "LocalAccountTokenFilterPolicy", 't': 4L,
+                         'value': struct.pack("<I", value)}
 
             h.node_set_value(key, new_value)
             h.commit(None)
@@ -435,7 +619,7 @@ class Windows(OSBase):
         addr = 'localhost'
         runas = '--runas=%s' % user
         winexe = subprocess.Popen(
-            ['winexe', '-U', user, "//%s" % addr, runas, command],
+            ['winexe', '-U', user, runas, "//%s" % addr, command],
             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
         stdout, stderr = winexe.communicate()