Revision f165adc0 image_creator/os_type/__init__.py

b/image_creator/os_type/__init__.py
31 31
# interpreted as representing official policies, either expressed
32 32
# or implied, of GRNET S.A.
33 33

  
34
from image_creator.util import output
34
from image_creator.util import output, FatalError
35 35

  
36
import textwrap
36 37
import re
37 38

  
38 39

  
......
43 44
    return wrapper
44 45

  
45 46

  
46
def exclude_task(func):
47
    func.excluded = True
48
    return func
47
def sysprep(enabled=True):
48
    def wrapper(func):
49
        func.sysprep = True
50
        func.enabled = enabled
51
        return func
52
    return wrapper
49 53

  
50 54

  
51 55
class OSBase(object):
52 56
    """Basic operating system class"""
57

  
53 58
    def __init__(self, rootdev, ghandler):
54 59
        self.root = rootdev
55 60
        self.g = ghandler
56 61

  
57
        self.sysprep_regexp = re.compile('^sysprep_')
58
        self.data_cleanup_regexp = re.compile('^data_cleanup_')
62
    def _is_sysprep(self, obj):
63
        return getattr(obj, 'sysprep', False) and callable(obj)
59 64

  
60
    def _print_task(self, task):
61
        name = task.__name__
65
    def list_syspreps(self):
62 66

  
63
        if self.sysprep_regexp.match(name):
64
            name = self.sysprep_regexp.sub("", name)
65
        elif self.data_cleanup_regexp.match(name):
66
            name = self.data_cleanup_regexp.sub("", name)
67
        else:
68
            raise FatalError("%s is not a task" % name)
67
        objs = [getattr(self, name) for name in dir(self) \
68
            if not name.startswith('_')]
69

  
70
        enabled = [x for x in objs if self._is_sysprep(x) and x.enabled]
71
        disabled = [x for x in objs if self._is_sysprep(x) and not x.enabled]
72

  
73
        return enabled, disabled
74

  
75
    def _sysprep_change_status(self, name, status):
76

  
77
        error_msg = "Syprep operation %s does not exist for %s" % \
78
                (name, self.__class__.__name__)
79

  
80
        method_name = name.replace('-', '_')
81
        method = None
82
        try:
83
            method = getattr(self, method_name)
84
        except AttributeError:
85
            raise FatalError(error_msg)
86

  
87
        if not self._is_sysprep(method):
88
            raise FatalError(error_msg)
89

  
90
        setattr(method.im_func, 'enabled', status)
91

  
92
    def enable_sysprep(self, name):
93
        """Enable a system preperation operation"""
94
        self._sysprep_change_status(name, True)
95

  
96
    def disable_sysprep(self, name):
97
        """Disable a system preperation operation"""
98
        self._sysprep_change_status(name, False)
99

  
100
    def print_syspreps(self):
101
        """Print enabled and disabled system preperation operations."""
69 102

  
70
        name = name.replace('_', '-')
103
        enabled, disabled = self.list_syspreps()
71 104

  
72
        output("  %s:\n    %s" % (name, task.__doc__))
105
        wrapper = textwrap.TextWrapper()
106
        wrapper.subsequent_indent = '\t'
107
        wrapper.initial_indent = '\t'
108

  
109
        output("Enabled system preperation operations:")
110
        if len(enabled) == 0:
111
            output("(none)")
112
        else:
113
            for sysprep in enabled:
114
                name = sysprep.__name__.replace('_', '-')
115
                descr = wrapper.fill(sysprep.__doc__)
116
                output('    %s:\n%s\n' % (name, descr))
117

  
118
        output("Disabled system preperation operations:")
119
        if len(disabled) == 0:
120
            output("(none)")
121
        else:
122
            for sysprep in disabled:
123
                name = sysprep.__name__.replace('_', '-')
124
                descr = wrapper.fill(sysprep.__doc__)
125
                output('    %s:\n%s\n' % (name, descr))
73 126

  
74 127
    @add_prefix
75 128
    def ls(self, directory):
......
133 186

  
134 187
        return meta
135 188

  
136
    def list_sysprep(self):
137
        """List all sysprep actions"""
138

  
139
        is_sysprep = lambda x: x.startswith('sysprep_') and \
140
                                                    callable(getattr(self, x))
141
        tasks = [getattr(self, x) for x in dir(self) if is_sysprep(x)]
142

  
143
        included = [t for t in tasks if not getattr(t, "excluded", False)]
144
        excluded = [t for t in tasks if getattr(t, "excluded", False)]
145

  
146
        return included, excluded
147

  
148
    def list_data_cleanup(self):
149
        """List all data_cleanup actions"""
150

  
151
        is_cleanup = lambda x: x.startswith('data_cleanup_') and \
152
                                                    callable(getattr(self, x))
153
        tasks = [getattr(self, x) for x in dir(self) if is_cleanup(x)]
154

  
155
        included = [t for t in tasks if not getattr(t, "excluded", False)]
156
        excluded = [t for t in tasks if getattr(t, "excluded", False)]
157

  
158
        return included, excluded
159

  
160
    def data_cleanup(self):
161
        """Cleanup sensitive data out of the OS image."""
162

  
163
        output('Cleaning up sensitive data out of the OS image:')
164

  
165
        tasks, _ = self.list_data_cleanup()
166
        size = len(tasks)
167
        cnt = 0
168
        for task in tasks:
169
            cnt += 1
170
            output(('(%d/%d)' % (cnt, size)).ljust(7), False)
171
            task()
172
        output()
173

  
174
    def sysprep(self):
189
    def do_sysprep(self):
175 190
        """Prepere system for image creation."""
176 191

  
177 192
        output('Preparing system for image creation:')
......
185 200
            task()
186 201
        output()
187 202

  
188
    def print_task(self, task):
189
        name = task.__name__
190

  
191
        if self.sysprep_regexp.match(name):
192
            name = self.sysprep_regexp.sub("", name)
193
        elif self.data_cleanup_regexp.match(name):
194
            name = self.data_cleanup_regexp.sub("", name)
195
        else:
196
            raise FatalError("%s is not a task" % name)
197

  
198
        name = name.replace('_', '-')
199

  
200
        output("  %s:\n    %s" % (name, task.__doc__))
201

  
202
    def print_data_cleanup(self):
203
        included, excluded = self.list_data_cleanup()
204

  
205
        output("Included data cleanup operations:")
206
        if len(included) == 0:
207
            ouput("(none)")
208
        else:
209
            for task in included:
210
                self._print_task(task)
211
        output("Ommited data cleanup operations:")
212
        if len(excluded) == 0:
213
            ouput("(none)")
214
        else:
215
            for task in excluded:
216
                self._print_task(task)
217

  
218
    def print_sysprep(self):
219
        included, excluded = self.list_sysprep()
220

  
221
        output("Included sysprep operations:")
222
        if len(included) == 0:
223
            ouput("(none)")
224
        else:
225
            for task in included:
226
                self._print_task(task)
227
        output("Ommited sysprep operations:")
228
        if len(excluded) == 0:
229
            output("(none)")
230
        else:
231
            for task in excluded:
232
                self._print_task(task)
233

  
234 203
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

Also available in: Unified diff