code.grnet.gr Git - snf-image-creator/blob - image_creator/dialog_main.py

   1 #!/usr/bin/env python
   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 import dialog
  37 import sys
  38 import os
  39 import textwrap
  40 import signal
  41 import StringIO
  42
  43 from image_creator import __version__ as version
  44 from image_creator.util import FatalError, MD5
  45 from image_creator.output.dialog import InitializationOutput, GaugeOutput
  46 from image_creator.disk import Disk
  47 from image_creator.os_type import os_cls
  48 from image_creator.kamaki_wrapper import Kamaki, ClientError
  49
  50 MSGBOX_WIDTH = 60
  51 YESNO_WIDTH = 50
  52 MENU_WIDTH = 70
  53 INPUTBOX_WIDTH=70
  54
  55
  56 class Reset(Exception):
  57     pass
  58
  59
  60 def confirm_exit(d, msg=''):
  61     return not d.yesno("%s Do you want to exit?" % msg, width=YESNO_WIDTH)
  62
  63
  64 def confirm_reset(d):
  65     return not d.yesno(
  66         "Are you sure you want to reset everything?",
  67         width=YESNO_WIDTH)
  68
  69 def extract_image(session):
  70     d = session['dialog']
  71     dir = os.getcwd()
  72     while 1:
  73         if dir and dir[-1] != os.sep:
  74             dir = dir + os.sep
  75
  76         (code, path) = d.fselect(dir, 10, 50, title="Save image as...")
  77         if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):
  78             return False
  79
  80         if os.path.isdir(path):
  81             dir=path
  82             continue
  83
  84         if os.path.isdir("%s.meta" % path):
  85             d.msgbox("Can't overwrite directory `%s.meta'" % path,
  86                      width=MSGBOX_WIDTH)
  87             continue
  88
  89         if os.path.isdir("%s.md5sum" % path):
  90             d.msgbox("Can't overwrite directory `%s.md5sum'" % path,
  91                      width=MSGBOX_WIDTH)
  92             continue
  93
  94         basedir = os.path.dirname(path)
  95         name = os.path.basename(path)
  96         if not os.path.exists(basedir):
  97             d.msgbox("Directory `%s' does not exist" % basedir,
  98                      width=MSGBOX_WIDTH)
  99             continue
 100
 101         dir = basedir
 102         if len(name) == 0:
 103             continue
 104
 105         files = ["%s%s" % (path, ext) for ext in ('', '.meta', '.md5sum')]
 106         overwrite = filter(os.path.exists, files)
 107
 108         if len(overwrite) > 0:
 109             if d.yesno("The following file(s) exist:\n"
 110                         "%s\nDo you want to overwrite them?" %
 111                         "\n".join(overwrite), width=YESNO_WIDTH):
 112                 continue
 113
 114         out = GaugeOutput(d, "Image Extraction", "Extracting image...")
 115         try:
 116             dev = session['device']
 117             if "checksum" not in session:
 118                 size = dev.meta['SIZE']
 119                 md5 = MD5(out)
 120                 session['checksum'] = md5.compute(session['snapshot'], size)
 121
 122             # Extract image file
 123             dev.out = out
 124             dev.dump(path)
 125
 126             # Extract metadata file
 127             out.output("Extracting metadata file...")
 128             metastring = '\n'.join(
 129                 ['%s=%s' % (k, v) for (k, v) in session['metadata'].items()])
 130             metastring += '\n'
 131             with open('%s.meta' % path, 'w') as f:
 132                 f.write(metastring)
 133             out.success('done')
 134
 135             # Extract md5sum file
 136             out.output("Extracting md5sum file...")
 137             md5str = "%s %s\n" % (session['checksum'], name)
 138             with open('%s.md5sum' % path, 'w') as f:
 139                 f.write(md5str)
 140             out.success("done")
 141
 142         finally:
 143             out.cleanup()
 144         d.msgbox("Image file `%s' was successfully extracted!" % path,
 145                  width=MSGBOX_WIDTH)
 146         break
 147
 148     return True
 149
 150
 151 def upload_image(session):
 152     d = session["dialog"]
 153     size = session['device'].meta['SIZE']
 154
 155     if "account" not in session:
 156         d.msgbox("You need to provide your ~okeanos login username before you "
 157                  "can upload images to pithos+", width=MSGBOX_WIDTH)
 158         return False
 159
 160     if "token" not in session:
 161         d.msgbox("You need to provide your ~okeanos account authentication "
 162                  "token before you can upload images to pithos+",
 163                  width=MSGBOX_WIDTH)
 164         return False
 165
 166     while 1:
 167         init=session["upload"] if "upload" in session else ''
 168         (code, answer) = d.inputbox("Please provide a filename:", init=init,
 169                                     width=INPUTBOX_WIDTH)
 170
 171         if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):
 172             return False
 173
 174         filename = answer.strip()
 175         if len(filename) == 0:
 176             d.msgbox("Filename cannot be empty", width=MSGBOX_WIDTH)
 177             continue
 178
 179         break
 180
 181     out = GaugeOutput(d, "Image Upload", "Uploading...")
 182     if 'checksum' not in session:
 183         md5 = MD5(out)
 184         session['checksum'] = md5.compute(session['snapshot'], size)
 185     try:
 186         kamaki = Kamaki(session['account'], session['token'], out)
 187         try:
 188             # Upload image file
 189             with open(session['snapshot'], 'rb') as f:
 190                 session["upload"] = kamaki.upload(f, size, filename,
 191                                                   "Calculating block hashes",
 192                                                   "Uploading missing blocks")
 193             # Upload metadata file
 194             out.output("Uploading metadata file...")
 195             metastring = '\n'.join(
 196                 ['%s=%s' % (k, v) for (k, v) in session['metadata'].items()])
 197             metastring += '\n'
 198             kamaki.upload(StringIO.StringIO(metastring), size=len(metastring),
 199                           remote_path="%s.meta" % filename)
 200             out.success("done")
 201
 202             # Upload md5sum file
 203             out.output("Uploading md5sum file...")
 204             md5str = "%s %s\n" % (session['checksum'], filename)
 205             kamaki.upload(StringIO.StringIO(md5str), size=len(md5str),
 206                           remote_path="%s.md5sum" % filename)
 207             out.success("done")
 208
 209         except ClientError as e:
 210             d.msgbox("Error in pithos+ client: %s" % e.message,
 211                      title="Pithos+ Client Error", width=MSGBOX_WIDTH)
 212             if 'upload' in session:
 213                 del session['upload']
 214             return False
 215     finally:
 216         out.cleanup()
 217
 218     d.msgbox("Image file `%s' was successfully uploaded to pithos+" % filename,
 219              width=MSGBOX_WIDTH)
 220     return True
 221
 222
 223 def register_image(session):
 224     d = session["dialog"]
 225
 226     if "account" not in session:
 227         d.msgbox("You need to provide your ~okeanos login username before you "
 228                  "can register an images to cyclades",
 229                  width=MSGBOX_WIDTH)
 230         return False
 231
 232     if "token" not in session:
 233         d.msgbox("You need to provide your ~okeanos account authentication "
 234                  "token before you can register an images to cyclades",
 235                  width=MSGBOX_WIDTH)
 236         return False
 237
 238     if "upload" not in session:
 239         d.msgbox("You need to have an image uploaded to pithos+ before you "
 240                  "can register it to cyclades",
 241                  width=MSGBOX_WIDTH)
 242         return False
 243
 244     while 1:
 245         (code, answer) = d.inputbox("Please provide a registration name:"
 246                                 " be registered:", width=INPUTBOX_WIDTH)
 247         if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):
 248             return False
 249
 250         name = answer.strip()
 251         if len(name) == 0:
 252             d.msgbox("Registration name cannot be empty", width=MSGBOX_WIDTH)
 253             continue
 254         break
 255
 256     out = GaugeOutput(d, "Image Registration", "Registrating image...")
 257     try:
 258         out.output("Registring image to cyclades...")
 259         try:
 260             kamaki = Kamaki(session['account'], session['token'], out)
 261             kamaki.register(name, session['upload'], session['metadata'])
 262             out.success('done')
 263         except ClientError as e:
 264             d.msgbox("Error in pithos+ client: %s" % e.message)
 265             return False
 266     finally:
 267         out.cleanup()
 268
 269     d.msgbox("Image `%s' was successfully registered to cyclades as `%s'" %
 270              (session['upload'], name), width=MSGBOX_WIDTH)
 271     return True
 272
 273
 274 def kamaki_menu(session):
 275     d = session['dialog']
 276     default_item = "Account"
 277     while 1:
 278         account = session["account"] if "account" in session else "<none>"
 279         token = session["token"] if "token" in session else "<none>"
 280         upload = session["upload"] if "upload" in session else "<none>"
 281         (code, choice) = d.menu(
 282             "Choose one of the following or press <Back> to go back.",
 283             width=MENU_WIDTH,
 284             choices=[("Account", "Change your ~okeanos username: %s" %
 285                       account),
 286                      ("Token", "Change your ~okeanos token: %s" % token),
 287                      ("Upload", "Upload image to pithos+"),
 288                      ("Register", "Register image to cyclades: %s" % upload)],
 289             cancel="Back",
 290             default_item=default_item,
 291             title="Image Registration Menu")
 292
 293         if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):
 294             return False
 295
 296         if choice == "Account":
 297             default_item = "Account"
 298             (code, answer) = d.inputbox(
 299                 "Please provide your ~okeanos account e-mail address:",
 300                 init=session["account"] if "account" in session else '',
 301                 width=70)
 302             if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):
 303                 continue
 304             if len(answer) == 0 and "account" in session:
 305                     del session["account"]
 306             else:
 307                 session["account"] = answer.strip()
 308                 default_item = "Token"
 309         elif choice == "Token":
 310             default_item = "Token"
 311             (code, answer) = d.inputbox(
 312                 "Please provide your ~okeanos account authetication token:",
 313                 init=session["token"] if "token" in session else '',
 314                 width=70)
 315             if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):
 316                 continue
 317             if len(answer) == 0 and "token" in session:
 318                 del session["token"]
 319             else:
 320                 session["token"] = answer.strip()
 321                 default_item = "Upload"
 322         elif choice == "Upload":
 323             if upload_image(session):
 324                 default_item = "Register"
 325             else:
 326                 default_item = "Upload"
 327         elif choice == "Register":
 328             if register_image(session):
 329                 return True
 330             else:
 331                 default_item = "Register"
 332
 333
 334 def main_menu(session):
 335     d = session['dialog']
 336     dev = session['device']
 337     d.setBackgroundTitle("OS: %s, Distro: %s" % (dev.ostype, dev.distro))
 338     actions = {"Register": kamaki_menu, "Extract": extract_image}
 339     default_item = "Customize"
 340
 341     while 1:
 342         (code, choice) = d.menu(
 343             "Choose one of the following or press <Exit> to exit.",
 344             width=MENU_WIDTH,
 345             choices=[("Customize", "Run various image customization tasks"),
 346                      ("Deploy", "Configure ~okeanos image deployment options"),
 347                      ("Register", "Register image to ~okeanos"),
 348                      ("Extract", "Dump image to local file system"),
 349                      ("Reset", "Reset everything and start over again"),
 350                      ("Help", "Get help for using snf-image-creator")],
 351             cancel="Exit",
 352             default_item=default_item,
 353             title="Image Creator for ~okeanos (snf-image-creator version %s)" %
 354                   version)
 355
 356         if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):
 357             if confirm_exit(d):
 358                 break
 359             else:
 360                 continue
 361
 362         if choice == "Reset":
 363             if confirm_reset(d):
 364                 d.infobox("Resetting snf-image-creator. Please wait...")
 365                 raise Reset
 366             else:
 367                 continue
 368         elif choice in actions:
 369             actions[choice](session)
 370
 371
 372 def select_file(d):
 373     root = os.sep
 374     while 1:
 375         (code, path) = d.fselect(root, 10, 50,
 376                                  title="Please select input media")
 377         if code in (d.DIALOG_CANCEL, d.DIALOG_ESC):
 378             if confirm_exit(d, "You canceled the media selection dialog box."):
 379                 sys.exit(0)
 380             else:
 381                 continue
 382
 383         if not os.path.exists(path):
 384             d.msgbox("The file you choose does not exist", width=MSGBOX_WIDTH)
 385             continue
 386         else:
 387             break
 388
 389     return path
 390
 391
 392 def collect_metadata(dev, out):
 393
 394     out.output("Collecting image metadata...")
 395     metadata = dev.meta
 396     dev.mount(readonly=True)
 397     cls = os_cls(dev.distro, dev.ostype)
 398     image_os = cls(dev.root, dev.g, out)
 399     dev.umount()
 400     metadata.update(image_os.meta)
 401     out.success("done")
 402
 403     return metadata
 404
 405
 406 def image_creator(d):
 407     basename = os.path.basename(sys.argv[0])
 408     usage = "Usage: %s [input_media]" % basename
 409     if len(sys.argv) > 2:
 410         sys.stderr.write("%s\n" % usage)
 411         return 1
 412
 413     if os.geteuid() != 0:
 414         raise FatalError("You must run %s as root" % basename)
 415
 416     media = sys.argv[1] if len(sys.argv) == 2 else select_file(d)
 417
 418     out = InitializationOutput(d)
 419     disk = Disk(media, out)
 420
 421     def signal_handler(signum, fram):
 422         out.cleanup()
 423         disk.cleanup()
 424
 425     signal.signal(signal.SIGINT, signal_handler)
 426     try:
 427         snapshot = disk.snapshot()
 428         dev = disk.get_device(snapshot)
 429
 430         metadata = collect_metadata(dev, out)
 431         out.cleanup()
 432
 433         # Make sure the signal handler does not call out.cleanup again
 434         def dummy(self):
 435             pass
 436         instancemethod = type(InitializationOutput.cleanup)
 437         out.cleanup = instancemethod(dummy, out, InitializationOutput)
 438
 439         session = {"dialog": d,
 440                    "disk": disk,
 441                    "snapshot": snapshot,
 442                    "device": dev,
 443                    "metadata": metadata}
 444
 445         main_menu(session)
 446         d.infobox("Thank you for using snf-image-creator. Bye", width=53)
 447     finally:
 448         disk.cleanup()
 449
 450     return 0
 451
 452
 453 def main():
 454
 455     d = dialog.Dialog(dialog="dialog")
 456
 457     while 1:
 458         try:
 459             try:
 460                 ret = image_creator(d)
 461                 sys.exit(ret)
 462             except FatalError as e:
 463                 msg = textwrap.fill(str(e), width=70)
 464                 d.infobox(msg, width=70, title="Fatal Error")
 465                 sys.exit(1)
 466         except Reset:
 467             continue
 468
 469 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :