Fix small bugs and clean up the code
[snf-image-creator] / image_creator / os_type / __init__.py
1 # -*- coding: utf-8 -*-
2 #
3 # Copyright 2012 GRNET S.A. All rights reserved.
4 #
5 # Redistribution and use in source and binary forms, with or
6 # without modification, are permitted provided that the following
7 # conditions are met:
8 #
9 #   1. Redistributions of source code must retain the above
10 #      copyright notice, this list of conditions and the following
11 #      disclaimer.
12 #
13 #   2. Redistributions in binary form must reproduce the above
14 #      copyright notice, this list of conditions and the following
15 #      disclaimer in the documentation and/or other materials
16 #      provided with the distribution.
17 #
18 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
19 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
20 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
21 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
22 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
25 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
26 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
27 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
28 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29 # POSSIBILITY OF SUCH DAMAGE.
30 #
31 # The views and conclusions contained in the software and
32 # documentation are those of the authors and should not be
33 # interpreted as representing official policies, either expressed
34 # or implied, of GRNET S.A.
35
36 """This package provides various classes for preparing different Operating
37 Systems for image creation.
38 """
39
40 from image_creator.util import FatalError
41
42 import textwrap
43 import re
44 from collections import namedtuple
45 from functools import wraps
46
47
48 def os_cls(distro, osfamily):
49     """Given the distro name and the osfamily, return the appropriate class"""
50     module = None
51     classname = None
52     try:
53         module = __import__("image_creator.os_type.%s" % distro,
54                             fromlist=['image_creator.os_type'])
55         classname = distro.capitalize()
56     except ImportError:
57         module = __import__("image_creator.os_type.%s" % osfamily,
58                             fromlist=['image_creator.os_type'])
59         classname = osfamily.capitalize()
60
61     return getattr(module, classname)
62
63
64 def add_prefix(target):
65     def wrapper(self, *args):
66         prefix = args[0]
67         return [prefix + path for path in target(self, *args)]
68     return wrapper
69
70
71 def sysprep(message, enabled=True, **kwargs):
72     """Decorator for system preparation tasks"""
73     def wrapper(method):
74         method.sysprep = True
75         method.enabled = enabled
76         method.executed = False
77
78         for key, val in kwargs.items():
79             setattr(method, key, val)
80
81         @wraps(method)
82         def inner(self, print_message=True):
83             if print_message:
84                 self.out.output(message)
85             return method(self)
86
87         return inner
88     return wrapper
89
90
91 def add_sysprep_param(name, default, description, validator=lambda x: True):
92     """Decorator for __init__ that adds the definition for a system preparation
93     parameter in an instance of a os_type class
94     """
95     def wrapper(init):
96         @wraps(init)
97         def inner(self, *args, **kwargs):
98             init(self, *args, **kwargs)
99             self.needed_sysprep_params[name] = \
100                 self.SysprepParam(default, description, validator)
101         return inner
102     return wrapper
103
104
105 def del_sysprep_param(name):
106     """Decorator for __init__ that deletes a previously added sysprep parameter
107     definition from an instance of a os_type class.
108     """
109     def wrapper(func):
110         @wraps(func)
111         def inner(self, *args, **kwargs):
112             del self.needed_sysprep_params[name]
113             func(self, *args, **kwargs)
114         return inner
115     return wrapper
116
117
118 class OSBase(object):
119     """Basic operating system class"""
120
121     SysprepParam = namedtuple('SysprepParam', 'default description validator')
122
123     def __init__(self, image, **kargs):
124         self.image = image
125
126         self.root = image.root
127         self.g = image.g
128         self.out = image.out
129
130         self.needed_sysprep_params = {}
131         self.sysprep_params = \
132             kargs['sysprep_params'] if 'sysprep_params' in kargs else {}
133
134         self.meta = {}
135         self.mounted = False
136
137     def collect_metadata(self):
138         """Collect metadata about the OS"""
139         try:
140             if not self.mount(readonly=True):
141                 raise FatalError("Unable to mount the media read-only")
142
143             self.out.output('Collecting image metadata ...', False)
144             self._do_collect_metadata()
145             self.out.success('done')
146         finally:
147             self.umount()
148
149         self.out.output()
150
151     def list_syspreps(self):
152         """Returns a list of sysprep objects"""
153         objs = [getattr(self, name) for name in dir(self)
154                 if not name.startswith('_')]
155
156         return [x for x in objs if self._is_sysprep(x) and x.executed is False]
157
158     def sysprep_info(self, obj):
159         """Returns information about a sysprep object"""
160         assert self._is_sysprep(obj), "Object is not a sysprep"
161
162         SysprepInfo = namedtuple("SysprepInfo", "name description")
163
164         return SysprepInfo(obj.__name__.replace('_', '-'),
165                            textwrap.dedent(obj.__doc__))
166
167     def get_sysprep_by_name(self, name):
168         """Returns the sysprep object with the given name"""
169         error_msg = "Syprep operation %s does not exist for %s" % \
170                     (name, self.__class__.__name__)
171
172         method_name = name.replace('-', '_')
173         method = None
174         try:
175             method = getattr(self, method_name)
176         except AttributeError:
177             raise FatalError(error_msg)
178
179         if not self._is_sysprep(method):
180             raise FatalError(error_msg)
181
182         return method
183
184     def enable_sysprep(self, obj):
185         """Enable a system preparation operation"""
186         setattr(obj.im_func, 'enabled', True)
187
188     def disable_sysprep(self, obj):
189         """Disable a system preparation operation"""
190         setattr(obj.im_func, 'enabled', False)
191
192     def print_syspreps(self):
193         """Print enabled and disabled system preparation operations."""
194
195         syspreps = self.list_syspreps()
196         enabled = [sysprep for sysprep in syspreps if sysprep.enabled]
197         disabled = [sysprep for sysprep in syspreps if not sysprep.enabled]
198
199         wrapper = textwrap.TextWrapper()
200         wrapper.subsequent_indent = '\t'
201         wrapper.initial_indent = '\t'
202         wrapper.width = 72
203
204         self.out.output("Enabled system preparation operations:")
205         if len(enabled) == 0:
206             self.out.output("(none)")
207         else:
208             for sysprep in enabled:
209                 name = sysprep.__name__.replace('_', '-')
210                 descr = wrapper.fill(textwrap.dedent(sysprep.__doc__))
211                 self.out.output('    %s:\n%s\n' % (name, descr))
212
213         self.out.output("Disabled system preparation operations:")
214         if len(disabled) == 0:
215             self.out.output("(none)")
216         else:
217             for sysprep in disabled:
218                 name = sysprep.__name__.replace('_', '-')
219                 descr = wrapper.fill(textwrap.dedent(sysprep.__doc__))
220                 self.out.output('    %s:\n%s\n' % (name, descr))
221
222     def print_sysprep_params(self):
223         """Print the system preparation parameter the user may use"""
224
225         self.out.output("Needed system preparation parameters:")
226
227         if len(self.needed_sysprep_params) == 0:
228             self.out.output("(none)")
229             return
230
231         for name, param in self.needed_sysprep_params.items():
232             self.out.output("\t%s (%s): %s" %
233                             (param.description, name,
234                              self.sysprep_params[name] if name in
235                              self.sysprep_params else "(none)"))
236
237     def do_sysprep(self):
238         """Prepare system for image creation."""
239
240         try:
241             if not self.mount(readonly=False):
242                 raise FatalError("Unable to mount the media read-write")
243
244             self.out.output('Preparing system for image creation:')
245
246             enabled = [task for task in self.list_syspreps() if task.enabled]
247
248             size = len(enabled)
249             cnt = 0
250             for task in enabled:
251                 cnt += 1
252                 self.out.output(('(%d/%d)' % (cnt, size)).ljust(7), False)
253                 task()
254                 setattr(task.im_func, 'executed', True)
255         finally:
256             self.umount()
257
258         self.out.output()
259
260     def mount(self, readonly=False):
261         """Mount image."""
262
263         if getattr(self, "mounted", False):
264             return True
265
266         mount_type = 'read-only' if readonly else 'read-write'
267         self.out.output("Mounting the media %s ..." % mount_type, False)
268
269         if not self._do_mount(readonly):
270             return False
271
272         self.mounted = True
273         self.out.success('done')
274         return True
275
276     def umount(self):
277         """Umount all mounted filesystems."""
278
279         self.out.output("Umounting the media ...", False)
280         self.g.umount_all()
281         self.mounted = False
282         self.out.success('done')
283
284     def _is_sysprep(self, obj):
285         """Checks if an object is a sysprep"""
286         return getattr(obj, 'sysprep', False) and callable(obj)
287
288     @add_prefix
289     def _ls(self, directory):
290         """List the name of all files under a directory"""
291         return self.g.ls(directory)
292
293     @add_prefix
294     def _find(self, directory):
295         """List the name of all files recursively under a directory"""
296         return self.g.find(directory)
297
298     def _foreach_file(self, directory, action, **kargs):
299         """Perform an action recursively on all files under a directory.
300
301         The following options are allowed:
302
303         * maxdepth: If defined the action will not be performed on
304           files that are below this level of directories under the
305           directory parameter.
306
307         * ftype: The action will only be performed on files of this
308           type. For a list of all allowed filetypes, see here:
309           http://libguestfs.org/guestfs.3.html#guestfs_readdir
310
311         * exclude: Exclude all files that follow this pattern.
312         """
313         maxdepth = None if 'maxdepth' not in kargs else kargs['maxdepth']
314         if maxdepth == 0:
315             return
316
317         # maxdepth -= 1
318         maxdepth = None if maxdepth is None else maxdepth - 1
319         kargs['maxdepth'] = maxdepth
320
321         exclude = None if 'exclude' not in kargs else kargs['exclude']
322         ftype = None if 'ftype' not in kargs else kargs['ftype']
323         has_ftype = lambda x, y: y is None and True or x['ftyp'] == y
324
325         for f in self.g.readdir(directory):
326             if f['name'] in ('.', '..'):
327                 continue
328
329             full_path = "%s/%s" % (directory, f['name'])
330
331             if exclude and re.match(exclude, full_path):
332                 continue
333
334             if has_ftype(f, 'd'):
335                 self._foreach_file(full_path, action, **kargs)
336
337             if has_ftype(f, ftype):
338                 action(full_path)
339
340     def _do_collect_metadata(self):
341         """helper method for collect_metadata"""
342         self.meta['ROOT_PARTITION'] = "%d" % self.g.part_to_partnum(self.root)
343         self.meta['OSFAMILY'] = self.g.inspect_get_type(self.root)
344         self.meta['OS'] = self.g.inspect_get_distro(self.root)
345         if self.meta['OS'] == "unknown":
346             self.meta['OS'] = self.meta['OSFAMILY']
347         self.meta['DESCRIPTION'] = self.g.inspect_get_product_name(self.root)
348
349     def _do_mount(self, readonly):
350         """helper method for mount"""
351         try:
352             self.g.mount_options('ro' if readonly else 'rw', self.root, '/')
353         except RuntimeError as msg:
354             self.out.warn("unable to mount the root partition: %s" % msg)
355             return False
356
357         return True
358
359 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :