Revision e8b1b48b

b/image_creator/disk.py
168 168
        self.bootable = bootable
169 169
        self.progress_bar = None
170 170
        self.guestfs_device = None
171
        self.size = None
172
        self.parttype = None
171
        self.meta = {}
173 172

  
174 173
        self.g = guestfs.GuestFS()
175 174
        self.g.add_drive_opts(self.real_device, readonly=0)
......
203 202
                            "We only support images with one filesystem.")
204 203
        self.root = roots[0]
205 204
        self.guestfs_device = self.g.part_to_dev(self.root)
206
        self.size = self.g.blockdev_getsize64(self.guestfs_device)
207
        self.parttype = self.g.part_get_parttype(self.guestfs_device)
205
        self.meta['SIZE'] = self.g.blockdev_getsize64(self.guestfs_device)
206
        self.meta['PARTITION_TABLE'] = \
207
                                self.g.part_get_parttype(self.guestfs_device)
208 208

  
209 209
        self.ostype = self.g.inspect_get_type(self.root)
210 210
        self.distro = self.g.inspect_get_distro(self.root)
......
253 253
        """Umount all mounted filesystems."""
254 254
        self.g.umount_all()
255 255

  
256
    def _last_partition(self):
257
        if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
258
            msg = "Unsupported partition table: %s. Only msdos and gpt " \
259
            "partition tables are supported" % self.meta['PARTITION_TABLE']
260
            raise FatalError(msg)
261

  
262
        is_extended = lambda p: self.g.part_get_mbr_id(
263
                                    self.guestfs_device, p['part_num']) == 5
264
        is_logical = lambda p: self.meta['PARTITION_TABLE'] != 'msdos' and \
265
                                                            p['part_num'] > 4
266

  
267
        partitions = self.g.part_list(self.guestfs_device)
268
        last_partition = partitions[-1]
269

  
270
        if is_logical(last_partition):
271
            # The disk contains extended and logical partitions....
272
            extended = [p for p in partitions if is_extended(p)][0]
273
            last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
274

  
275
            # check if extended is the last primary partition
276
            if last_primary['part_num'] > extended['part_num']:
277
                last_partition = last_primary
278

  
279
        return last_partition
280

  
256 281
    def shrink(self):
257 282
        """Shrink the disk.
258 283

  
......
262 287

  
263 288
        ATTENTION: make sure unmount is called before shrink
264 289
        """
265
        output("Shrinking image (this may take a while)...", False)
266

  
267
        if self.parttype not in 'msdos' 'gpt':
268
            raise FatalError("You have a %s partition table. "
269
                "Only msdos and gpt partitions are supported" % self.parttype)
290
        get_fstype = lambda p: self.g.vfs_type("%s%d" % \
291
                                        (self.guestfs_device, p['part_num']))
292
        is_logical = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
293
                                                            p['part_num'] > 4
294
        is_extended = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
295
                self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
296

  
297
        part_add = lambda ptype, start, stop: \
298
                    self.g.part_add(self.guestfs_device, ptype, start, stop)
299
        part_del = lambda p: self.g.part_del(self.guestfs_device, p)
300
        part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
301
        part_set_id = lambda p, id: self.g.part_set_mbr_id(
302
                                                    self.guestfs_device, p, id)
303
        part_get_bootable = lambda p: self.g.part_get_bootable(
304
                                                        self.guestfs_device, p)
305
        part_set_bootable = lambda p, bootable: self.g.part_set_bootable(
306
                                            self.guestfs_device, p, bootable)
307

  
308
        MB = 2 ** 20
270 309

  
271
        last_partition = self.g.part_list(self.guestfs_device)[-1]
272

  
273
        if self.parttype == 'msdos' and last_partition['part_num'] > 4:
274
            raise FatalError("This disk contains logical partitions. "
275
                                    "Only primary partitions are supported.")
276

  
277
        part_dev = "%s%d" % (self.guestfs_device, last_partition['part_num'])
278
        fs_type = self.g.vfs_type(part_dev)
279
        if not re.match("ext[234]", fs_type):
280
            warn("Don't know how to resize %s partitions." % fs_type)
281
            return self.size
310
        output("Shrinking image (this may take a while)...", False)
282 311

  
312
        last_part = None
313
        fstype = None
314
        while True:
315
            last_part = self._last_partition()
316
            fstype = get_fstype(last_part)
317

  
318
            if fstype == 'swap':
319
                self.meta['SWAP'] = "%d:%s" % \
320
                        (last_part['part_num'],
321
                        (last_part['part_size'] + MB - 1) // MB)
322
                part_del(last_part['part_num'])
323
                continue
324
            elif is_extended(last_part):
325
                part_del(last_part['part_num'])
326
                continue
327

  
328
            self.meta['SIZE'] = last_part['part_end'] + 1
329
            break
330

  
331
        if not re.match("ext[234]", fstype):
332
            warn("Don't know how to resize %s partitions." % fstype)
333
            return self.meta['SIZE']
334

  
335
        part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
283 336
        self.g.e2fsck_f(part_dev)
284 337
        self.g.resize2fs_M(part_dev)
285 338

  
......
290 343
            filter(lambda x: x[0] == 'Block count', out)[0][1])
291 344

  
292 345
        sector_size = self.g.blockdev_getss(self.guestfs_device)
293

  
294
        start = last_partition['part_start'] / sector_size
346
        start = last_part['part_start'] / sector_size
295 347
        end = start + (block_size * block_cnt) / sector_size - 1
296 348

  
297
        self.g.part_del(self.guestfs_device, last_partition['part_num'])
298
        self.g.part_add(self.guestfs_device, 'p', start, end)
349
        if is_logical(last_part):
350
            partitions = self.g.part_list(self.guestfs_device)
351

  
352
            logical = []  # logical partitions
353
            for partition in partitions:
354
                if partition['part_num'] < 4:
355
                    continue
356
                logical.append({
357
                    'num': partition['part_num'],
358
                    'start': partition['part_start'] / sector_size,
359
                    'end': partition['part_end'] / sector_size,
360
                    'id': part_get_(partition['part_num']),
361
                    'bootable': part_get_bootable(partition['part_num'])
362
                })
363

  
364
            logical[-1]['end'] = end  # new end after resize
365

  
366
            # Recreate the extended partition
367
            extended = [p for p in partitions if self._is_extended(p)][0]
368
            part_del(extended['part_num'])
369
            part_add('e', extended['part_start'], end)
370

  
371
            # Create all the logical partitions back
372
            for l in logical:
373
                part_add('l', l['start'], l['end'])
374
                part_set_id(l['num'], l['id'])
375
                part_set_bootable(l['num'], l['bootable'])
376
        else:
377
            # Recreate the last partition
378
            if self.meta['PARTITION_TABLE'] == 'msdos':
379
                last_part['id'] = part_get_id(last_part['part_num'])
380

  
381
            last_part['bootable'] = part_get_bootable(last_part['part_num'])
382
            part_del(last_part['part_num'])
383
            part_add('p', start, end)
384
            part_set_bootable(last_part['part_num'], last_part['bootable'])
299 385

  
300
        self.size = (end + 1) * sector_size
301
        success("new size is %dMB" % ((self.size + 2 ** 20 - 1) // 2 ** 20))
386
            if self.meta['PARTITION_TABLE'] == 'msdos':
387
                part_set_id(last_part['part_num'], last_part['id'])
302 388

  
303
        if self.parttype == 'gpt':
389
        new_size = (end + 1) * sector_size
390
        success("new size is %dMB" % ((new_size + MB - 1) // MB))
391

  
392
        if self.meta['PARTITION_TABLE'] == 'gpt':
304 393
            ptable = GPTPartitionTable(self.real_device)
305
            self.size = ptable.shrink(self.size)
394
            self.meta['SIZE'] = ptable.shrink(new_size)
395
        else:
396
            self.meta['SIZE'] = new_size
306 397

  
307
        return self.size
398
        return self.meta['SIZE']
308 399

  
309 400
    def dump(self, outfile):
310 401
        """Dumps the content of device into a file.
......
312 403
        This method will only dump the actual payload, found by reading the
313 404
        partition table. Empty space in the end of the device will be ignored.
314 405
        """
315
        blocksize = 2 ** 22  # 4MB
316
        progress_size = (self.size + 2 ** 20 - 1) // 2 ** 20  # in MB
406
        MB = 2 ** 20
407
        blocksize = 4 * MB  # 4MB
408
        size = self.meta['SIZE']
409
        progress_size = (size + MB - 1) // MB  # in MB
317 410
        progressbar = progress("Dumping image file: ", 'mb')
318 411
        progressbar.max = progress_size
319 412

  
320 413
        with open(self.real_device, 'r') as src:
321 414
            with open(outfile, "w") as dst:
322
                left = self.size
415
                left = size
323 416
                offset = 0
324 417
                progressbar.next()
325 418
                while left > 0:
......
327 420
                    sent = sendfile(dst.fileno(), src.fileno(), offset, length)
328 421
                    offset += sent
329 422
                    left -= sent
330
                    progressbar.goto((self.size - left) // 2 ** 20)
423
                    progressbar.goto((size - left) // MB)
331 424
        output("\rDumping image file...\033[K", False)
332 425
        success('image file %s was successfully created' % outfile)
333 426

  
b/image_creator/main.py
192 192
        dev.umount()
193 193

  
194 194
        size = options.shrink and dev.shrink() or dev.size
195
        metadata['SIZE'] = str(size // 2 ** 20)
195
        metadata.update(dev.meta)
196 196

  
197 197
        checksum = md5(snapshot, size)
198 198

  

Also available in: Unified diff