Synnefo » snf-image-creator

import time

losetup = get_command('losetup')

mount = get_command('mount')

umount = get_command('umount')

MKFS_OPTS = {

    'ext2': ['-F'],

    'ext3': ['-F'],

    'ext4': ['-F'],

    'reiserfs': ['-ff'],

    'btrfs': [],

    'minix': [],

    'xfs': ['-f'],

    'jfs': ['-f'],

    'ntfs': ['-F'],

    'msdos': [],

    'vfat': []

    }

            root = findfs(root).stdout.strip()

        if not re.match('/dev/[hsv]d[a-z][1-9]*$', root):

            raise FatalError("Don't know how to handle root device: %s" % root)

        out.success(root)

        disk_file = re.split('[0-9]', root)[0]

        device = parted.Device(disk_file)

        self.disk = parted.Disk(device)

        FileSystemEntry = namedtuple('FileSystemEntry',

                                     'dev mpoint fs opts freq passno')

                yield FileSystemEntry(*entry)

    def _get_mount_options(self, device):

        return None

    def _create_partition_table(self, image):

        if self.disk.type != 'msdos':

        first_sector = self.disk.getPrimaryPartitions()[0].geometry.start

        dd('if=%s' % self.disk.device.path, 'of=%s' % image,

           'bs=%d' % self.disk.device.sectorSize,

        extended = self.disk.getExtendedPartition()

        logical = self.disk.getLogicalPartitions()

            dd('if=%s' % self.disk.device.path, 'of=%s' % image,

    def _get_partitions(self, disk):

        Partition = namedtuple('Partition', 'num start end type fs')

        for p in disk.partitions:

            num = p.number

            start = p.geometry.start

            end = p.geometry.end

            ptype = p.type

            fs = p.fileSystem.type if p.fileSystem is not None else ''

            partitions.append(Partition(num, start, end, ptype, fs))

        return partitions

    def _shrink_partitions(self, image):

        new_end = self.disk.device.getLength()

        image_dev = parted.Device(image)

        partitions = self._get_partitions(self.disk)

            size = (last.end - last.start + 1) * self.disk.device.sectorSize

        mount_options = self._get_mount_options(

                self.disk.getPartitionBySector(last.start).path)

        if mount_options is not None: 

            stat = os.statvfs(mount_options.mpoint)

            # Shrink the last partition. The new size should be the size of the

            # occupied blocks

            new_size = (blcks * stat.f_frsize) // self.disk.device.sectorSize

            # Parted may have changed this for better alignment

            part_end = image_disk.getPartitionBySector(last.start).geometry.end

            last = last._replace(end=part_end)

            partitions[-1] = last

            # Leave 2048 blocks at the end.

            new_end = new_size + 2048

        return new_end

    def _map_partition(self, dev, num, start, end):

        name = os.path.basename(dev)

        tablefd, table = tempfile.mkstemp()

        try:

            size = end - start + 1

            os.write(tablefd, "0 %d linear %s %d" % (size, dev, start))

            dmsetup('create', "%sp%d" % (name, num), table)

        finally:

            os.unlink(table)

        return "/dev/mapper/%sp%d" % (name, num)

    def _unmap_partition(self, dev):

        if not os.path.exists(dev):

            return

        dmsetup('remove', dev.split('/dev/mapper/')[1])

        time.sleep(0.1)

    def _mount(self, target, devs):

        devs.sort(key=lambda d: d[1])

        for dev, mpoint in devs:

            absmpoint = os.path.abspath(target + mpoint)

            if not os.path.exists(absmpoint):

                os.makedirs(absmpoint)

            mount(dev, absmpoint)

    def _umount_all(self, target):

        mpoints = []

        for entry in self._read_fstable('/proc/mounts'):

            if entry.mpoint.startswith(os.path.abspath(target)):

                    mpoints.append(entry.mpoint)

        mpoints.sort()

        for mpoint in reversed(mpoints):

            umount(mpoint)

    def _create_filesystems(self, image):

        partitions = self._get_partitions(parted.Disk(parted.Device(image)))

        filesystems = {}

        for p in self.disk.partitions:

            filesystems[p.number] = self._get_mount_options(p.path)

        unmounted = filter(lambda p: filesystems[p.num] is None, partitions)

        mounted = filter(lambda p: filesystems[p.num] is not None, partitions)

        # For partitions that are not mounted right now, we can simply dd them

        # into the image.

        for p in unmounted:

            dd('if=%s' % self.disk.device.path, 'of=%s' % image,

               'count=%d' % (p.end - p.start + 1), 'conv=notrunc',

               'seek=%d' % p.start, 'skip=%d' % p.start)

        loop = str(losetup('-f', '--show', image)).strip()

        mapped = {}

        try:

            for p in mounted:

                i =  p.num

                mapped[i] = self._map_partition(loop, i, p.start, p.end)

            # Create the file systems

            for i, dev in mapped.iteritems():

                fs = filesystems[i].fs

                self.out.output('Creating %s filesystem on partition %d ... ' %

                    (fs, i), False)

                get_command('mkfs.%s' % fs)(*(MKFS_OPTS[fs] + [dev]))

                self.out.success('done')

            target = tempfile.mkdtemp()

            try:

                absmpoints = self._mount(target,

                    [(mapped[i], filesystems[i].mpoint) for i in mapped.keys()]

                )

            finally:

                self._umount_all(target)

                os.rmdir(target)

        finally:

            for dev in mapped.values():

                self._unmap_partition(dev)

            losetup('-d', loop)

    def create_image(self):

        image = '/mnt/%s.diskdump' % uuid.uuid4().hex

        disk_size = self.disk.device.getLength() * self.disk.device.sectorSize

        truncate("-s", "%d" % disk_size, image)

        self._create_partition_table(image)

        end_sector = self._shrink_partitions(image)

        # Check if the available space is enough to host the image

        dirname = os.path.dirname(image)

        size = (end_sector + 1) * self.disk.device.sectorSize

        self.out.output("Examining available space in %s ..." % dirname, False)

        stat = os.statvfs(dirname)

        available = stat.f_bavail * stat.f_frsize

        if available <= size:

            raise FatalError('Not enough space in %s to host the image' %

                             dirname)

        self.out.success("sufficient")

        self._create_filesystems(image)

        return image