Use libguestfs to access the image
[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 class DiskError(Exception): pass
17
18 class Disk(object):
19
20     def __init__(self, source):
21         self._cleanup_jobs = []
22         self._devices = []
23         self.source = source
24
25     def _add_cleanup(self, job, *args):
26         self._cleanup_jobs.append((job, args))
27
28     def _losetup(self, fname):
29         loop = losetup.find_unused_loop_device()
30         loop.mount(fname)
31         self._add_cleanup(loop.unmount)
32         return loop.device
33
34     def _dir_to_disk(self):
35         raise NotImplementedError
36
37     def cleanup(self):
38         while len(self._devices):
39             device = self._devices.pop()
40             device.destroy()
41             
42         while len(self._cleanup_jobs):
43             job, args = self._cleanup_jobs.pop()
44             job(*args)
45
46     def get_device(self):
47         sourcedev = self.source
48         mode = os.stat(self.source).st_mode
49         if stat.S_ISDIR(mode):
50             return self._losetup(self._dir_to_disk())
51         elif stat.S_ISREG(mode):
52             sourcedev = self._losetup(self.source)
53         elif not stat.S_ISBLK(mode):
54             raise ValueError("Value for self.source is invalid")
55
56         # Take a snapshot and return it to the user
57         size = blockdev('--getsize', sourcedev)
58         cowfd, cow = tempfile.mkstemp()
59         self._add_cleanup(os.unlink, cow)
60         # Create 1G cow file
61         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k' ,'seek=%d' % (1024*1024))
62         cowdev = self._losetup(cow)
63
64         snapshot = uuid.uuid4().hex
65         tablefd, table = tempfile.mkstemp()
66         try:
67             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
68                                         (int(size), sourcedev, cowdev))
69             dmsetup('create', snapshot, table)
70             self._add_cleanup(dmsetup, 'remove', snapshot)
71         finally:
72             os.unlink(table)
73
74         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
75         self._devices.append(new_device)
76         return new_device
77
78     def destroy_device(self, device):
79         self._devices.remove(device)
80         device.destroy()
81
82 class DiskDevice(object):
83
84     def __init__(self, device, bootable = True):
85         self.device = device
86         self.bootable = bootable
87
88         self.g = guestfs.GuestFS()
89         self.g.add_drive_opts(device, readonly = 0)
90         self.g.launch()
91         roots = self.g.inspect_os()
92         if len(roots) == 0:
93             raise DiskError("No operating system found")
94         if len(roots) > 1:
95             raise DiskError("Multiple operating systems found")
96
97         self.root = roots[0]
98     
99     def destroy(self):
100         self.g.umount_all()
101         self.g.sync()
102         # Close the guestfs handler
103         del self.g
104     
105     def get_image_metadata(self):
106         meta = {}
107         meta["OSFAMILY"] = self.g.inspect_get_type(self.root)
108         meta["OS"] = self.g.inspect_get_distro(self.root)
109         meta["description"] = self.g.inspect_get_product_name(self.root)
110         return meta
111
112     def mount(self):
113         mps = g.inspect_get_mountpoints(self.root)
114         # Sort the keys to mount the fs in a correct order.
115         # / should be mounted befor /boot, etc
116         def compare (a, b):
117             if len(a[0]) > len(b[0]): return 1
118             elif len(a[0]) == len(b[0]): return 0
119             else: return -1
120         mps.sort(compare)
121         for mp, dev in mps:
122             try:
123                 self.g.mount(dev, mp)
124             except RuntimeError as msg:
125                 print "%s (ignored)" % msg
126
127 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :