add setuptools support

7 files changed, 309 insertions, 8 deletions

diff --git a/vmdebootstrap/__init__.py b/vmdebootstrap/__init__.py
index e69de29..8b13789 100644
--- a/vmdebootstrap/__init__.py
+++ b/vmdebootstrap/__init__.py
@@ -0,0 +1 @@
+
diff --git a/vmdebootstrap/base.py b/vmdebootstrap/base.py
index c2799e7..1843d56 100644
--- a/vmdebootstrap/base.py
+++ b/vmdebootstrap/base.py
@@ -44,6 +44,7 @@ def runcmd(argv, stdin='', ignore_fail=False, env=None, **kwargs):
     return out
 
 
+# FIXME: use contextmanager
 def mount_wrapper(rootdir):
     runcmd(['mount', '/dev', '-t', 'devfs', '-obind',
             '%s' % os.path.join(rootdir, 'dev')])
@@ -131,3 +132,45 @@ class Base(object):
             else:
                 create_user(userpass)
                 delete_password(rootdir, userpass)
+
+    def customize(self, rootdir):
+        script = self.settings['customize']
+        if not script:
+            return
+        if not os.path.exists(script):
+            example = os.path.join("/usr/share/vmdebootstrap/examples/", script)
+            if not os.path.exists(example):
+                self.message("Unable to find %s" % script)
+                return
+            script = example
+        self.message('Running customize script %s' % script)
+        logging.info("rootdir=%s image=%s", rootdir, self.settings['image'])
+        logging.debug(
+            "%s usage: %s", self.settings['image'],
+            runcmd(['du', self.settings['image']]))
+        with open('/dev/tty', 'w') as tty:
+            try:
+                cliapp.runcmd([script, rootdir, self.settings['image']], stdout=tty, stderr=tty)
+            except IOError:
+                subprocess.call([script, rootdir, self.settings['image']])
+        logging.debug(
+            "%s usage: %s", self.settings['image'],
+            runcmd(['du', self.settings['image']]))
+
+    def optimize_image(self, rootdir):
+        """
+        Filing up the image with zeros will increase its compression rate
+        """
+        if not self.settings['sparse']:
+            zeros = os.path.join(rootdir, 'ZEROS')
+            self.runcmd_unchecked(['dd', 'if=/dev/zero', 'of=' + zeros, 'bs=1M'])
+            runcmd(['rm', '-f', zeros])
+
+    def append_serial_console(self, rootdir):
+        if self.settings['serial-console']:
+            serial_command = self.settings['serial-console-command']
+            logging.debug('adding getty to serial console')
+            inittab = os.path.join(rootdir, 'etc/inittab')
+            # to autologin, serial_command can contain '-a root'
+            with open(inittab, 'a') as ftab:
+                ftab.write('\nS0:23:respawn:%s\n' % serial_command)
diff --git a/vmdebootstrap/codenames.py b/vmdebootstrap/codenames.py
index 045d5b2..25cd94d 100644
--- a/vmdebootstrap/codenames.py
+++ b/vmdebootstrap/codenames.py
@@ -35,6 +35,7 @@ class Codenames(Base):
         super(Codenames, self).__init__()
         self.debian_info = DebianDistroInfo()
         self.ubuntu_info = UbuntuDistroInfo()
+        self.settings = None
 
     def define_settings(self, settings):
         self.settings = settings
@@ -58,4 +59,3 @@ class Codenames(Base):
         if not self.debian_info.valid(suite):
             return False
         return suite == self.debian_info.stable(limit)
-
diff --git a/vmdebootstrap/extlinux.py b/vmdebootstrap/extlinux.py
index 7286282..d5861a9 100644
--- a/vmdebootstrap/extlinux.py
+++ b/vmdebootstrap/extlinux.py
@@ -1,4 +1,6 @@
-#!/usr/bin/env python
+"""
+  Wrapper for Extlinux support
+"""
 # -*- coding: utf-8 -*-
 #
 #  extlinux.py
@@ -25,6 +27,8 @@ import cliapp
 import logging
 from vmdebootstrap.base import Base, runcmd
 
+# pylint: disable=missing-docstring
+
 
 class ExtLinux(Base):
 
@@ -88,3 +92,14 @@ append initrd=%(initrd)s root=UUID=%(uuid)s ro %(kserial)s
         runcmd(['extlinux', '--install', rootdir])
         runcmd(['sync'])
         time.sleep(2)
+
+    def install_mbr(self):
+        if not self.settings['mbr'] and not self.settings['extlinux']:
+            return
+        if os.path.exists("/sbin/install-mbr"):
+            self.message('Installing MBR')
+            runcmd(['install-mbr', self.settings['image']])
+        else:
+            msg = "mbr enabled but /sbin/install-mbr not found" \
+                  " - please install the mbr package."
+            raise cliapp.AppException(msg)
diff --git a/vmdebootstrap/filesystem.py b/vmdebootstrap/filesystem.py
index e0a16d9..1ccaeed 100644
--- a/vmdebootstrap/filesystem.py
+++ b/vmdebootstrap/filesystem.py
@@ -20,7 +20,14 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from vmdebootstrap.base import Base
+
+import os
+import cliapp
+import logging
+import subprocess
+from vmdebootstrap.base import Base, runcmd
+
+# pylint: disable=missing-docstring
 
 
 class Filesystem(Base):
@@ -29,3 +36,227 @@ class Filesystem(Base):
 
     def __init__(self):
         super(Filesystem, self).__init__()
+        self.settings = None
+        self.devices = {
+            'rootdir': None,
+            'rootdev': None,
+            'bootdev': None,
+            'boottype': None,
+            'roottype': None,
+            'swapdev': None,
+        }
+
+    def define_settings(self, settings):
+        self.settings = settings
+        self.devices['roottype'] = self.settings['roottype']
+
+    def chown(self):
+        if not self.settings['owner']:
+            return
+        # Change image owner after completed build
+        if self.settings['image']:
+            filename = self.settings['image']
+        elif self.settings['tarball']:
+            filename = self.settings['tarball']
+        else:
+            return
+        self.message("Changing owner to %s" % self.settings["owner"])
+        subprocess.call(["chown", self.settings["owner"], filename])
+
+    def update_initramfs(self):
+        rootdir = self.devices['rootdir']
+        if not rootdir:
+            raise cliapp.AppException("rootdir not set")
+        cmd = os.path.join('usr', 'sbin', 'update-initramfs')
+        if os.path.exists(os.path.join(str(rootdir), cmd)):
+            self.message("Updating the initramfs")
+            runcmd(['chroot', rootdir, cmd, '-u'])
+
+    def setup_kpartx(self):
+        bootindex = None
+        swapindex = None
+        out = runcmd(['kpartx', '-avs', self.settings['image']])
+        if self.settings['bootsize'] and self.settings['swap'] > 0:
+            bootindex = 0
+            rootindex = 1
+            swapindex = 2
+            parts = 3
+        elif self.settings['use-uefi']:
+            bootindex = 0
+            rootindex = 1
+            parts = 2
+        elif self.settings['use-uefi'] and self.settings['swap'] > 0:
+            bootindex = 0
+            rootindex = 1
+            swapindex = 2
+            parts = 3
+        elif self.settings['bootsize']:
+            bootindex = 0
+            rootindex = 1
+            parts = 2
+        elif self.settings['swap'] > 0:
+            rootindex = 0
+            swapindex = 1
+            parts = 2
+        else:
+            rootindex = 0
+            parts = 1
+        boot = None
+        swap = None
+        devices = [line.split()[2]
+                   for line in out.splitlines()
+                   if line.startswith('add map ')]
+        if len(devices) != parts:
+            msg = 'Surprising number of partitions - check output of losetup -a'
+            logging.debug("%s", runcmd(['losetup', '-a']))
+            logging.debug("%s: devices=%s parts=%s", msg, devices, parts)
+            raise cliapp.AppException(msg)
+        root = '/dev/mapper/%s' % devices[rootindex]
+        if self.settings['bootsize'] or self.settings['use-uefi']:
+            boot = '/dev/mapper/%s' % devices[bootindex]
+        if self.settings['swap'] > 0:
+            swap = '/dev/mapper/%s' % devices[swapindex]
+        self.devices['rootdev'] = root
+        self.devices['bootdev'] = boot
+        self.devices['swap'] = swap
+
+    def mkfs(self, device, fstype):
+        self.message('Creating filesystem %s' % fstype)
+        runcmd(['mkfs', '-t', fstype, device])
+
+    def create_fstab(self):
+        rootdir = self.devices['rootdir']
+        rootdev = self.devices['rootdev']
+        bootdev = self.devices['bootdev']
+        boottype = self.devices['boottype']
+
+        def fsuuid(device):
+            out = runcmd(['blkid', '-c', '/dev/null', '-o', 'value',
+                          '-s', 'UUID', device])
+            return out.splitlines()[0].strip()
+
+        if rootdev:
+            rootdevstr = 'UUID=%s' % fsuuid(rootdev)
+        else:
+            rootdevstr = '/dev/sda1'
+
+        if bootdev and not self.settings['use-uefi']:
+            bootdevstr = 'UUID=%s' % fsuuid(bootdev)
+        else:
+            bootdevstr = None
+
+        if not rootdir:
+            raise cliapp.AppException("rootdir not set")
+
+        fstab = os.path.join(str(rootdir), 'etc', 'fstab')
+        with open(fstab, 'w') as fstab:
+            fstab.write('proc /proc proc defaults 0 0\n')
+            fstab.write('%s / %s errors=remount-ro 0 1\n' % (rootdevstr, roottype))
+            if bootdevstr:
+                fstab.write('%s /boot %s errors=remount-ro 0 2\n' % (bootdevstr, boottype))
+                if self.settings['swap'] > 0:
+                    fstab.write("/dev/sda3 swap swap defaults 0 0\n")
+            elif self.settings['swap'] > 0:
+                fstab.write("/dev/sda2 swap swap defaults 0 0\n")
+
+    def squash(self):
+        """
+        Run squashfs on the image.
+        """
+        if not self.settings['squash']:
+            return
+        if not os.path.exists('/usr/bin/mksquashfs'):
+            logging.warning("Squash selected but mksquashfs not found!")
+            return
+        logging.debug(
+            "%s usage: %s", self.settings['image'],
+            runcmd(['du', self.settings['image']]))
+        self.message("Running mksquashfs")
+        suffixed = "%s.squashfs" % self.settings['image']
+        if os.path.exists(suffixed):
+            os.unlink(suffixed)
+        msg = runcmd(
+            ['mksquashfs', self.settings['image'],
+             suffixed,
+             '-no-progress', '-comp', 'xz'], ignore_fail=False)
+        logging.debug(msg)
+        check_size = os.path.getsize(suffixed)
+        if check_size < (1024 * 1024):
+            logging.warning(
+                "%s appears to be too small! %s bytes",
+                suffixed, check_size)
+        else:
+            logging.debug("squashed size: %s", check_size)
+            os.unlink(self.settings['image'])
+            self.settings['image'] = suffixed
+            logging.debug(
+                "%s usage: %s", self.settings['image'],
+                runcmd(['du', self.settings['image']]))
+
+    def configure_apt(self):
+        rootdir = self.devices['rootdir']
+        if not self.settings['configure-apt'] or not self.settings['apt-mirror']:
+            return
+        if not rootdir:
+            raise cliapp.AppException("rootdir not set")
+        # use the distribution and mirror to create an apt source
+        self.message("Configuring apt to use distribution and mirror")
+        conf = os.path.join(str(rootdir), 'etc', 'apt', 'sources.list.d', 'base.list')
+        logging.debug('configure apt %s', conf)
+        mirror = self.settings['mirror']
+        if self.settings['apt-mirror']:
+            mirror = self.settings['apt-mirror']
+            self.message("Setting apt mirror to %s" % mirror)
+        os.unlink(os.path.join(str(rootdir), 'etc', 'apt', 'sources.list'))
+        source = open(conf, 'w')
+        line = 'deb %s %s main\n' % (mirror, self.settings['distribution'])
+        source.write(line)
+        line = '#deb-src %s %s main\n' % (mirror, self.settings['distribution'])
+        source.write(line)
+        source.close()
+        # ensure the apt sources have valid lists
+        runcmd(['chroot', rootdir, 'apt-get', '-qq', 'update'])
+
+    def list_installed_pkgs(self):
+        if not self.settings['pkglist']:
+            return
+        rootdir = self.devices['rootdir']
+        # output the list of installed packages for sources identification
+        self.message("Creating a list of installed binary package names")
+        out = runcmd(['chroot', rootdir,
+                      'dpkg-query', '-W', "-f='${Package}.deb\n'"])
+        with open('dpkg.list', 'w') as dpkg:
+            dpkg.write(out)
+
+    def remove_udev_persistent_rules(self):
+        rootdir = self.devices['rootdir']
+        if not rootdir:
+            raise cliapp.AppException("rootdir not set")
+        self.message('Removing udev persistent cd and net rules')
+        for x in ['70-persistent-cd.rules', '70-persistent-net.rules']:
+            pathname = os.path.join(str(rootdir), 'etc', 'udev', 'rules.d', x)
+            if os.path.exists(pathname):
+                logging.debug('rm %s', pathname)
+                os.remove(pathname)
+            else:
+                logging.debug('not removing non-existent %s', pathname)
+
+    def set_hostname(self):
+        rootdir = self.devices['rootdir']
+        hostname = self.settings['hostname']
+        if not rootdir:
+            raise cliapp.AppException("rootdir not set")
+        with open(os.path.join(str(rootdir), 'etc', 'hostname'), 'w') as f:
+            f.write('%s\n' % hostname)
+
+        etc_hosts = os.path.join(str(rootdir), 'etc', 'hosts')
+        try:
+            with open(etc_hosts, 'r') as f:
+                data = f.read()
+            with open(etc_hosts, 'w') as f:
+                for line in data.splitlines():
+                    if line.startswith('127.0.0.1'):
+                        line += ' %s' % hostname
+                    f.write('%s\n' % line)
+        except IOError:
+            pass
diff --git a/vmdebootstrap/grub.py b/vmdebootstrap/grub.py
index 8d55ce1..22d7658 100644
--- a/vmdebootstrap/grub.py
+++ b/vmdebootstrap/grub.py
@@ -69,8 +69,10 @@ class GrubHandler(Base):
         except cliapp.AppException as exc:
             logging.warning(exc)
             self.message("Failed. Is grub2-common installed? Using extlinux.")
-            self.install_extlinux(rootdev, rootdir)  # FIXME!
+            umount_wrapper(rootdir)
+            return False
         umount_wrapper(rootdir)
+        return True
 
     def install_grub_uefi(self, rootdir):
         self.message("Configuring grub-uefi")
@@ -86,8 +88,9 @@ class GrubHandler(Base):
             self.message(
                 "Failed to configure grub-uefi for %s" %
                 self.settings['arch'])
-            umount_wrapper(rootdir)
-        self.uefi.configure_efi()  # FIXME
+        umount_wrapper(rootdir)
+
+    def install_extra_grub_uefi(self, rootdir):
         extra = str(arch_table[self.settings['arch']]['extra'])
         if extra:
             target = arch_table[extra]['target']
@@ -99,5 +102,4 @@ class GrubHandler(Base):
                 logging.warning(exc)
                 self.message(
                     "Failed to configure grub-uefi for %s" % extra)
-            self.uefi.configure_extra_efi()  # FIXME
         umount_wrapper(rootdir)
diff --git a/vmdebootstrap/uefi.py b/vmdebootstrap/uefi.py
index 82bf1eb..1acae14 100644
--- a/vmdebootstrap/uefi.py
+++ b/vmdebootstrap/uefi.py
@@ -24,7 +24,12 @@
 import os
 import cliapp
 import logging
-from vmdebootstrap.base import Base, runcmd
+from vmdebootstrap.base import (
+    Base,
+    runcmd,
+    mount_wrapper,
+    umount_wrapper,
+)
 
 # pylint: disable=missing-docstring
 
@@ -106,18 +111,22 @@ class Uefi(Base):
         so needs to be after grub and kernel already installed.
         """
         self.message('Configuring EFI')
+        mount_wrapper()
         efi_removable = str(arch_table[self.settings['arch']]['removable'])
         efi_install = str(arch_table[self.settings['arch']]['install'])
         self.message('Installing UEFI support binary')
         self.copy_efi_binary(efi_removable, efi_install)
+        umount_wrapper()
 
     def configure_extra_efi(self):
         extra = str(arch_table[self.settings['arch']]['extra'])
         if extra:
+            mount_wrapper()
             efi_removable = str(arch_table[extra]['removable'])
             efi_install = str(arch_table[extra]['install'])
             self.message('Copying UEFI support binary for %s' % extra)
             self.copy_efi_binary(efi_removable, efi_install)
+            umount_wrapper()
 
     def partition_esp(self):
         espsize = self.settings['esp-size'] / (1024 * 1024)