Make DiskDevice.shrink() return the new disk size
[snf-image-creator] / image_creator / disk.py
1 #!/usr/bin/env python
2
3 import losetup
4 import stat
5 import os
6 import tempfile
7 import uuid
8 import re
9 import sys
10 import guestfs
11
12 from pbs import dmsetup
13 from pbs import blockdev
14 from pbs import dd
15
16
17 class DiskError(Exception):
18     pass
19
20
21 class Disk(object):
22
23     def __init__(self, source):
24         self._cleanup_jobs = []
25         self._devices = []
26         self.source = source
27
28     def _add_cleanup(self, job, *args):
29         self._cleanup_jobs.append((job, args))
30
31     def _losetup(self, fname):
32         loop = losetup.find_unused_loop_device()
33         loop.mount(fname)
34         self._add_cleanup(loop.unmount)
35         return loop.device
36
37     def _dir_to_disk(self):
38         raise NotImplementedError
39
40     def cleanup(self):
41         while len(self._devices):
42             device = self._devices.pop()
43             device.destroy()
44
45         while len(self._cleanup_jobs):
46             job, args = self._cleanup_jobs.pop()
47             job(*args)
48
49     def get_device(self):
50         sourcedev = self.source
51         mode = os.stat(self.source).st_mode
52         if stat.S_ISDIR(mode):
53             return self._losetup(self._dir_to_disk())
54         elif stat.S_ISREG(mode):
55             sourcedev = self._losetup(self.source)
56         elif not stat.S_ISBLK(mode):
57             raise ValueError("Value for self.source is invalid")
58
59         # Take a snapshot and return it to the user
60         size = blockdev('--getsize', sourcedev)
61         cowfd, cow = tempfile.mkstemp()
62         self._add_cleanup(os.unlink, cow)
63         # Create 1G cow sparse file
64         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
65         cowdev = self._losetup(cow)
66
67         snapshot = uuid.uuid4().hex
68         tablefd, table = tempfile.mkstemp()
69         try:
70             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
71                                         (int(size), sourcedev, cowdev))
72             dmsetup('create', snapshot, table)
73             self._add_cleanup(dmsetup, 'remove', snapshot)
74         finally:
75             os.unlink(table)
76
77         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
78         self._devices.append(new_device)
79         return new_device
80
81     def destroy_device(self, device):
82         self._devices.remove(device)
83         device.destroy()
84
85
86 class DiskDevice(object):
87
88     def __init__(self, device, bootable=True):
89         self.device = device
90         self.bootable = bootable
91
92         self.g = guestfs.GuestFS()
93
94         self.g.set_trace(1)
95
96         self.g.add_drive_opts(device, readonly=0)
97         self.g.launch()
98         roots = self.g.inspect_os()
99         if len(roots) == 0:
100             raise DiskError("No operating system found")
101         if len(roots) > 1:
102             raise DiskError("Multiple operating systems found")
103
104         self.root = roots[0]
105         self.ostype = self.g.inspect_get_type(self.root)
106         self.distro = self.g.inspect_get_distro(self.root)
107
108     def destroy(self):
109         self.g.umount_all()
110         self.g.sync()
111         # Close the guestfs handler
112         self.g.close()
113         del self.g
114
115     def mount(self):
116         mps = self.g.inspect_get_mountpoints(self.root)
117
118         # Sort the keys to mount the fs in a correct order.
119         # / should be mounted befor /boot, etc
120         def compare(a, b):
121             if len(a[0]) > len(b[0]):
122                 return 1
123             elif len(a[0]) == len(b[0]):
124                 return 0
125             else:
126                 return -1
127         mps.sort(compare)
128         for mp, dev in mps:
129             try:
130                 self.g.mount(dev, mp)
131             except RuntimeError as msg:
132                 print "%s (ignored)" % msg
133
134     def umount(self):
135         self.g.umount_all()
136
137     def shrink(self):
138         dev = self.g.part_to_dev(self.root)
139         parttype = self.g.part_get_parttype(dev)
140         if parttype != 'msdos':
141             raise DiskError("You have a %s partition table. "
142                 "Only msdos partitions are supported" % parttype)
143
144         last_partition = self.g.part_list(dev)[-1]
145
146         if last_partition['part_num'] > 4:
147             raise DiskError("This disk contains logical partitions. "
148                 "Only primary partitions are supported.")
149
150         part_dev = "%s%d" % (dev, last_partition['part_num'])
151         fs_type = self.g.vfs_type(part_dev)
152         if not re.match("ext[234]", fs_type):
153             print "Warning, don't know how to resize %s partitions" % vfs_type
154             return
155
156         self.g.e2fsck_f(part_dev)
157         self.g.resize2fs_M(part_dev)
158         output = self.g.tune2fs_l(part_dev)
159         block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
160         block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
161
162         sector_size = self.g.blockdev_getss(dev)
163
164         start = last_partition['part_start'] / sector_size
165         end = start + (block_size * block_cnt) / sector_size - 1
166
167         return (end + 1) * sector_size
168
169
170
171 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :