Revision fbdf1d8f image_creator/dialog_wizard.py

b/image_creator/dialog_wizard.py
34 34
# or implied, of GRNET S.A.
35 35

  
36 36
import dialog
37
import time
38
import StringIO
37 39

  
38
from image_creator.kamaki_wrapper import Kamaki
40
from image_creator.kamaki_wrapper import Kamaki, ClientError
41
from image_creator.util import MD5, FatalError
42
from image_creator.output.cli import OutputWthProgress
39 43

  
40 44
PAGE_WIDTH = 70
41 45

  
......
65 69
class WizardPage:
66 70
    NEXT = 1
67 71
    PREV = -1
72
    EXIT = -255
73

  
74
    def run(self, session, index, total):
75
        raise NotImplementedError
76

  
77

  
78
class WizardInputPage(WizardPage):
68 79

  
69 80
    def __init__(self, name, message, **kargs):
70 81
        self.name = name
......
96 107
        return self.NEXT
97 108

  
98 109

  
110
class WizardYesNoPage(WizardPage):
111

  
112
    def __init__(self, message, **kargs):
113
        self.message = message
114
        self.title = kargs['title'] if 'title' in kargs else ''
115

  
116
    def run(self, session, index, total):
117
        d = session['dialog']
118

  
119
        while True:
120
            ret = d.yesno(self.message, width=PAGE_WIDTH, ok_label="Yes",
121
                    cancel="Back", extra_button=1, extra_label="Quit",
122
                    title="(%d/%d) %s" % (index + 1, total, self.title))
123

  
124
            if ret == d.DIALOG_CANCEL:
125
                return self.PREV
126
            elif ret == d.DIALOG_EXTRA:
127
                return self.EXIT
128
            elif ret == d.DIALOG_OK:
129
                return self.NEXT
130

  
131

  
99 132
def wizard(session):
100 133

  
101
    name = WizardPage("ImageName", "Please provide a name for the image:",
134
    name = WizardInputPage("ImageName", "Please provide a name for the image:",
102 135
                      title="Image Name", init=session['device'].distro)
103
    descr = WizardPage("ImageDescription",
136
    descr = WizardInputPage("ImageDescription",
104 137
        "Please provide a description for the image:",
105 138
        title="Image Description", empty=True,
106 139
        init=session['metadata']['DESCRIPTION'] if 'DESCRIPTION' in
107 140
        session['metadata'] else '')
108
    account = WizardPage("account",
141
    account = WizardInputPage("account",
109 142
        "Please provide your ~okeanos account e-mail:",
110 143
        title="~okeanos account information", init=Kamaki.get_account())
111
    token = WizardPage("token",
144
    token = WizardInputPage("token",
112 145
        "Please provide your ~okeanos account token:",
113 146
        title="~okeanos account token", init=Kamaki.get_token())
114 147

  
148
    msg = "Do you wish to continue with the image extraction process?"
149
    proceed = WizardYesNoPage(msg, title="Confirmation")
150

  
115 151
    w = Wizard(session)
152

  
116 153
    w.add_page(name)
117 154
    w.add_page(descr)
118 155
    w.add_page(account)
119 156
    w.add_page(token)
120

  
121
    return w.run()
157
    w.add_page(proceed)
158

  
159
    if w.run():
160
        extract_image(session)
161
    else:
162
        return False
163

  
164
    return True
165

  
166

  
167
def extract_image(session):
168
    disk = session['disk']
169
    device = session['device']
170
    snapshot = session['snapshot']
171
    image_os = session['image_os']
172
    wizard = session['wizard']
173

  
174
    out = OutputWthProgress(True)
175
    #Initialize the output
176
    disk.out = out
177
    device.out = out
178
    image_os.out = out
179

  
180
    out.output()
181

  
182
    #Sysprep
183
    device.mount(False)
184
    image_os.do_sysprep()
185
    metadata = image_os.meta
186
    device.umount()
187

  
188
    #Shrink
189
    size = device.shrink()
190

  
191
    #MD5
192
    md5 = MD5(out)
193
    checksum = md5.compute(snapshot, size)
194

  
195
    #Metadata
196
    metastring = '\n'.join(
197
        ['%s=%s' % (key, value) for (key, value) in metadata.items()])
198
    metastring += '\n'
199

  
200
    out.output()
201
    try:
202
        out.output("Uploading image to pithos:")
203
        kamaki = Kamaki(wizard['account'], wizard['token'], out)
204

  
205
        name = "%s-%s.diskdump" % (wizard['ImageName'],
206
                                   time.strftime("%Y%m%d%H%M"))
207
        pithos_file = ""
208
        with open(snapshot, 'rb') as f:
209
            pithos_file = kamaki.upload(f, size, name,
210
                                         "(1/4)  Calculating block hashes",
211
                                         "(2/4)  Uploading missing blocks")
212

  
213
        out.output("(3/4)  Uploading metadata file...", False)
214
        kamaki.upload(StringIO.StringIO(metastring), size=len(metastring),
215
                      remote_path="%s.%s" % (name, 'meta'))
216
        out.success('done')
217
        out.output("(4/4)  Uploading md5sum file...", False)
218
        md5sumstr = '%s %s\n' % (checksum, name)
219
        kamaki.upload(StringIO.StringIO(md5sumstr), size=len(md5sumstr),
220
                      remote_path="%s.%s" % (name, 'md5sum'))
221
        out.success('done')
222
        out.output()
223

  
224
        out.output('Registring image to ~okeanos...', False)
225
        kamaki.register(wizard['ImageName'], pithos_file, metadata)
226
        out.success('done')
227
        out.output()
228
    except ClientError as e:
229
        raise FatalError("Pithos client: %d %s" % (e.status, e.message))
122 230

  
123 231
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :

Also available in: Unified diff