mic2 0.24.12-1 / usr / bin / mic-image-writer

#!/usr/bin/python -tt
#
# mic-image-writer : write an image to usb disk or read to a file
#                       from usb disk
#
# Copyright 2009, Intel Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

import os
import subprocess
import time
import signal
import sys
import gettext
import optparse
import select
import dbus
import re
import mic.utils.misc as misc

try:
    from dbus.mainloop.glib import DBusGMainLoop
except:
    raise

_ = gettext.lgettext
COLOR_BLACK = "\033[00m"
COLOR_RED =   "\033[1;31m"
COLOR_BLUE =  "\033[1;34m"


def errmsg(msg):
    print >> sys.stderr, _("%s%s%s" % (COLOR_RED, msg, COLOR_BLACK))
    sys.exit(1)

def parse_options(args):
    parser = optparse.OptionParser("Usage: %prog [options] [image file]")
    parser.add_option("-c", "--console", action="store_true", dest="console",
                      default=False, help="Run in console mode")
    parser.add_option("-g", "--gui", action="store_true", dest="gui",
                      default=False, help="Run in GUI mode")
    (options, args) = parser.parse_args()
    if len(args) > 1:
        errmsg("Error: too much arguments.")
    elif len(args) == 1:
        if not os.path.isfile(args[0]):
            errmsg("Error: invalid image file.")
        options.image_file = os.path.abspath(os.path.expanduser(args[0]))
    else:
        options.image_file = None
    return options

class DBus_udisk:
    def __init__(self):
        dbus_loop = DBusGMainLoop(set_as_default=True)
        self.usbdisks = {}
        self.bus = dbus.SystemBus(mainloop = dbus_loop)
        self.UDISKS_IFACE = "org.freedesktop.UDisks"
        self.udisks = self.bus.get_object(self.UDISKS_IFACE, "/org/freedesktop/UDisks")

    def connect_signals(self, callback_add, callback_remove):
        # Connect with the Added and Removed signals
        self.udisks.connect_to_signal('DeviceAdded', callback_add)
        self.udisks.connect_to_signal('DeviceRemoved', callback_remove)

    def get_usb_list(self):
        ''' Dictionary drives store the info of usb disks with device file path as the key,
            such as '/dev/sdb'. The element is also dict with following format:
            { 'vendor': 'String of Vendor',
              'model': 'String of Model',
              'partition': list(partitions device object of this drive) }
        '''
        drives = {}
        devices = self.udisks.EnumerateDevices(dbus_interface = self.UDISKS_IFACE)
        for device in devices:
            try:
                dev_obj = self.bus.get_object(self.UDISKS_IFACE, device)
                dev_iface_p = dbus.Interface(dev_obj, 'org.freedesktop.DBus.Properties')

                if (not dev_iface_p.Get(self.UDISKS_IFACE, "DeviceIsSystemInternal")
                    and dev_iface_p.Get(self.UDISKS_IFACE, 'DeviceIsDrive')):
                    if dev_iface_p.Get(self.UDISKS_IFACE, 'DriveConnectionInterface') == 'usb':
                        dev_file = str(dev_iface_p.Get(self.UDISKS_IFACE, 'DeviceFile'))

                        if dev_file not in drives.keys() :
                            drives[dev_file] = {}

                        drives[dev_file]['vendor'] = dev_iface_p.Get(self.UDISKS_IFACE, 'DriveVendor')
                        drives[dev_file]['model'] = dev_iface_p.Get(self.UDISKS_IFACE, 'DriveModel')

                    elif dev_iface_p.Get(self.UDISKS_IFACE, 'DeviceIsPartition'):
                        partition = dev_iface_p.Get(self.UDISKS_IFACE, 'DeviceFile')
                        m = re.match("(\D+)(\d+)", partition)
                        if m:
                            dev_file = m.group(1)

                            if dev_file not in drives.keys() :
                                drives[dev_file] = {}

                            if 'partition' not in drives[dev_file].keys():
                                drives[dev_file]['partition'] = []

                            ''' Keep the partition's info in list'''
                            drives[dev_file]['partition'].append(dev_obj)

            except dbus.DBusException, e:
                errmsg(e)
        self.usbdisks = drives
        return drives

    def unmount(self, device):
        usbdisk = self.usbdisks[device]

        if 'partition' in usbdisk.keys():
            for dev_obj in usbdisk['partition']:
                dev_iface_p = dbus.Interface(dev_obj, 'org.freedesktop.DBus.Properties')
                dev_iface = dbus.Interface(dev_obj, 'org.freedesktop.UDisks.Device')
                try:
                    if dev_iface_p.Get(self.UDISKS_IFACE, 'DeviceIsMounted'):
                        dev_iface.FilesystemUnmount('')
                except dbus.exceptions.DBusException, e:
                    errmsg(e)

class DBus_hal:

    def __init__(self):
        dbus_loop = DBusGMainLoop(set_as_default=True)
        bus = dbus.SystemBus(mainloop = dbus_loop)
        hal_obj = bus.get_object("org.freedesktop.Hal",
                                 "/org/freedesktop/Hal/Manager")
        hal = dbus.Interface(hal_obj, "org.freedesktop.Hal.Manager")
        """Initialize dbus"""
        self.bus = bus
        self.hal = hal

    def connect_signals(self, callback_add, callback_remove):
        # Connect with the Added and Removed signals
        self.hal.connect_to_signal('DeviceAdded', callback_add)
        self.hal.connect_to_signal('DeviceRemoved', callback_remove)

    def get_usb_list(self):
        drives = {}
        devices = []
        devices = self.hal.FindDeviceByCapability("storage")
        for device in devices:
            try:
                dev_obj = self.bus.get_object("org.freedesktop.Hal", device)
                dev = dbus.Interface(dev_obj, "org.freedesktop.Hal.Device")
                if dev.GetProperty("storage.bus") == "usb" and \
                    dev.GetProperty("storage.removable"):
                    if not dev.GetProperty("block.is_volume"):
                        tmpdevice = str(dev.GetProperty('block.device'))
                        tmpvendor = str(dev.GetProperty('info.vendor'))
                        tmpproduct = str(dev.GetProperty('info.product'))
                        drives[tmpdevice]={"vendor":tmpvendor, "model":tmpproduct}
            except dbus.DBusException, e:
                continue
        return drives

    def hal_unmount(self, device):
        dev_obj = self.bus.get_object("org.freedesktop.Hal", device)
        dev = dbus.Interface(dev_obj, "org.freedesktop.Hal.Device")
        if dev.GetProperty('volume.is_mounted'):
            try:
                dev_obj.Unmount([""], dbus_interface="org.freedesktop.Hal.Device.Volume")
            except dbus.DBusException, msg:
                device = str(dev.GetProperty('block.device'))
                raise IOError("Error: failed to unmount %s, %s" % (device, msg))

    def unmount(self, usbdisk):
        partitions = self.hal.FindDeviceStringMatch('block.device', usbdisk)

        for partition in partitions:
            dev_obj = self.bus.get_object("org.freedesktop.Hal", partition)
            dev = dbus.Interface(dev_obj, "org.freedesktop.Hal.Device")
            if dev.GetProperty("block.is_volume"):
                device = str(dev.GetProperty('block.device'))
                self.hal_unmount(partition)
            else: # iterate over children looking for a volume
                if dev.GetProperty("storage.bus") == "usb" and \
                    dev.GetProperty("storage.removable"):
                    children = self.hal.FindDeviceStringMatch("info.parent",
                                                              partition)
                    for child in children:
                        dev_obj = self.bus.get_object("org.freedesktop.Hal", child)
                        dev = dbus.Interface(dev_obj, "org.freedesktop.Hal.Device")
                        if dev.GetProperty("block.is_volume"):
                            self.hal_unmount(child)

class USBWriterWindow:
    def __init__(self, the_image_file, usbdevices):
        self.p = None
        self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
        self.window.set_resizable(False)

        self.window.connect("delete-event", self.close)
        self.window.set_title("MIC2 USB Writer")
        self.window.set_border_width(0)

        vbox = gtk.VBox(True, 3)
        vbox.set_border_width(5)
        self.window.add(vbox)

        self.image_file = gtk.FileChooserButton("Select an Image to write")
        if the_image_file:
            self.image_file.set_filename(the_image_file)
        self.image_file.set_size_request(200, -1)

        # Create a centering alignment object
        align = gtk.Alignment(0.50, 0.50, 0.9, 0.10)
        align.add(self.image_file)

        image_frame = gtk.Frame("Image File")
        image_frame.set_shadow_type(gtk.SHADOW_IN)
        image_frame.set_size_request(260, 60)
        image_frame.add(align)

        self.combobox_disks = gtk.combo_box_new_text()
        self.combobox_disks.set_size_request(80, -1)

        # Create a centering alignment object
        align = gtk.Alignment(0.50, 0.50, 0.9, 0.10)
        align.add(self.combobox_disks)

        usb_frame = gtk.Frame("USB Disk")
        usb_frame.set_shadow_type(gtk.SHADOW_IN)
        usb_frame.set_size_request(120, 60)
        usb_frame.add(align)

        hbox = gtk.HBox(False, 2)
        hbox.set_border_width(0)

        hbox.pack_start(image_frame, True, True, 0)
        hbox.pack_start(usb_frame, True, True, 0)
        vbox.pack_start(hbox, True, True, 0)

        # Create the ProgressBar
        self.pbar = gtk.ProgressBar()
        progress_frame = gtk.Frame("Progress")

        # Create a centering alignment object
        align = gtk.Alignment(0.50, 0.50, 0.9, 0.10)
        progress_frame.add(align)

        align.add(self.pbar)

        vbox.pack_start(progress_frame, True, True, 0)

        self.write_button = gtk.Button("_Write")
        self.read_button = gtk.Button("_Read")
        self.cancel_button = gtk.Button("_Cancel")
        self.exit_button = gtk.Button("E_xit")

        hbox = gtk.HBox(False, 4)
        hbox.pack_start(self.write_button, False, False, 5)
        hbox.pack_start(self.read_button, False, False, 5)
        hbox.pack_start(self.cancel_button, False, False, 5)
        hbox.pack_start(self.exit_button, False, False, 5)

        # Create a right-side alignment object
        align = gtk.Alignment(1.0, 0.5, 0.05, 0.20)
        vbox.pack_start(align, True, True, 0)

        align.add(hbox)

        self.cancel_button.connect("clicked", self.cancel_operation)
        self.write_button.connect("clicked", self.write_image)
        self.read_button.connect("clicked", self.read_image)
        self.exit_button.connect("clicked", self.exit)

        self.usbdevices = usbdevices
        self.usbdevices.connect_signals(self.populate_devices, self.populate_devices)
        usbdisks_list = self.usbdevices.get_usb_list()

        self.usb_disks_num = 0
        self.refresh_usblist(usbdisks_list)
        self.window.show_all()

    def close(self, widget, event):
        self.exit(widget, event)
        return True

    def exit(self, widget, data=None):
        if self.p:
            if self.p.poll() == None:
                os.kill(self.p.pid, signal.SIGTERM)
            while self.p.poll() == None:
                while gtk.events_pending():
                    gtk.main_iteration()
        gtk.main_quit()

    def cancel_operation(self, widget):
        if self.p and self.p.poll() == None:
            os.kill(self.p.pid, signal.SIGTERM)

    def populate_devices(self, *args, **kw):
        self.refresh_usblist(self.usbdevices.get_usb_list())
        pass

    def refresh_usblist(self, usblist):
        self.usb_disks = []
        if self.usb_disks_num > 0:
            for i in range(0, self.usb_disks_num):
                self.combobox_disks.remove_text(0)

        for key in usblist.keys():
            self.combobox_disks.append_text(key)
            self.usb_disks.append(key)

        self.usb_disks_num = len(usblist)
        if not self.combobox_disks.get_active() >= 0:
            self.combobox_disks.set_active(0)
        self.combobox_disks.show()

    def msgbox(self, msg, type = 0):
        dlgtypes = [gtk.MESSAGE_WARNING, gtk.MESSAGE_INFO, gtk.MESSAGE_ERROR]
        dlgtitles = ["Warning", "Info", "Error"]
        dlg = gtk.MessageDialog(None, 0, dlgtypes[type], gtk.BUTTONS_OK, msg)
        dlg.set_title(dlgtitles[type])
        dlg.set_default_size(100, 100)
        dlg.set_transient_for(self.window)
        dlg.set_keep_above(True)
        dlg.present()
        dlg.set_position (gtk.WIN_POS_CENTER_ON_PARENT)
        dlg.set_border_width(2)
        dlg.set_modal(True)
        dlg.run()
        dlg.hide()
        dlg.destroy()
        while gtk.events_pending():
            gtk.main_iteration()

    def __is_usbdisk(self, disk_no):
        if disk_no < 0 or not self.usb_disks:
            return False

        if self.usb_disks[disk_no]:
            # It's a usb disk
            return True
        return False

    def set_progress(self):
        percent = self.elapsed_time*1.0/self.total_time
        if percent > 1.0:
            percent = 1.0
        self.pbar.set_fraction(percent)
        self.pbar.set_text("estimated time: %ds, elapsed time: %ds, progress: %d%%" %
                           (self.total_time, self.elapsed_time, int(percent*100)))

    def disable_some_widgets(self):
        self.image_file.set_sensitive(False)
        self.combobox_disks.set_sensitive(False)
        self.write_button.set_sensitive(False)
        self.read_button.set_sensitive(False)

    def enable_some_widgets(self):
        self.image_file.set_sensitive(True)
        self.combobox_disks.set_sensitive(True)
        self.write_button.set_sensitive(True)
        self.read_button.set_sensitive(True)

    def write_image(self, widget):
        image_file = self.image_file.get_filename()
        if not image_file:
            self.msgbox("Please specify a image file.")
            return

        if not os.path.exists(image_file) or not os.path.isfile(image_file):
            self.msgbox("Please specify a valid image file.")

        usb_disk = self.combobox_disks.get_active()
        if not self.__is_usbdisk(usb_disk):
            self.msgbox("Please insert or select a usb disk.")
            return

        self.disable_some_widgets()
        argv = ["/bin/dd", "if=%s" % image_file, "bs=1M", "of=%s" % self.usb_disks[usb_disk]]

        try:
            self.usbdevices.unmount(self.usb_disks[usb_disk])
        except IOError, e:
            self.msgbox(str(e))
            self.enable_some_widgets()
            return
        self.p = subprocess.Popen(argv,
                             stdout = subprocess.PIPE,
                             stderr = subprocess.STDOUT,
                             stdin = subprocess.PIPE,
                             close_fds = True)

        image_size = os.path.getsize(image_file)
        image_size_mb = image_size/1024/1024
        self.total_time = image_size_mb/10
        self.elapsed_time = 0
        while self.p.poll() == None:
              time.sleep(1)
              self.elapsed_time += 1
              self.set_progress()
              while gtk.events_pending():
                 gtk.main_iteration()
        (stdout, stderr) = self.p.communicate()
        if self.p.returncode != 0:
            for line in stdout.split('\n'):
                print line
            self.msgbox("Failed!!!", 2)
        else:
            self.set_progress()
            self.msgbox("Done successfully!!!", 1)
        self.enable_some_widgets()

    def read_image(self, widget):
        pass


if __name__ == "__main__":
    misc.setlocale()
    options = parse_options(sys.argv[1:])
    if not options.console or options.gui:
        try:
            import pygtk
            pygtk.require('2.0')
            if not os.environ.has_key("DISPLAY") or not os.environ["DISPLAY"]:
                errmsg("Error: No env $DISPLAY (add 'Defaults env_keep += \"DISPLAY\"'to \"/etc/sudoers\")")
            import gtk, gobject
            options.gui = True
        except:
            if options.gui:
                errmsg("Error: failed to run in GUI mode")
            options.console = True

    if options.console:
        if not options.image_file:
            errmsg("Error: please provide image file.")

    if os.geteuid () != 0:
        errmsg("You must run mic-image-writer as root")

    try:
        usbdevices = DBus_udisk()
    except dbus.exceptions.DBusException:
        usbdevices = DBus_hal()

    if options.gui:
        USBWriterWindow(options.image_file, usbdevices)
        gtk.main()
    else:
        usb_list = usbdevices.get_usb_list()

        if not usb_list:
            print "Please insert your USB stick..."
        while not usb_list:
            time.sleep(2)
            usb_list = usbdevices.get_usb_list()

        index = 1
        device_index={}

        print "Available usb disk(s):"
        for key in usb_list.keys():
            print "\t[%d] %s %s %s" %(index, key, usb_list[key]["vendor"], usb_list[key]['model'])
            device_index[index]= key
            index += 1
        while True:
            choice = raw_input("Please choice [1..%d] ? " % (index - 1,))
            if choice.isdigit() and int(choice) > 0 and int(choice) < index:
                choice = int(choice)
                break
            else:
                print "Invalid choice"
        usbdevices.unmount(device_index[choice])
        image_size = os.path.getsize(options.image_file)
        image_size_mb = image_size/1024/1024
        total_time = image_size_mb/10
        elapsed_time = 0
        argv = ["/bin/dd", "if=%s" % options.image_file, "bs=1M", "of=%s" % device_index[choice]]
        print "Source: %s" % options.image_file
        print "Target: %s" % device_index[choice]
        print "Image size: %d MB" % image_size_mb
        print "Estimated time: %d seconds" % total_time
        p = subprocess.Popen(argv,
                             stdout = subprocess.PIPE,
                             stderr = subprocess.STDOUT,
                             stdin = subprocess.PIPE,
                             close_fds = True)
        while p.poll() == None:
              time.sleep(1)
              elapsed_time += 1
              percent = elapsed_time*100/total_time
              if percent > 100:
                  percent = 100
              print _("\r%sElapsed time: %d, \tprogress: %d%%%s") % (COLOR_BLUE, elapsed_time, percent, COLOR_BLACK),
              sys.stdout.flush()
        if p.returncode == 0:
            print _("\r%sElapsed time: %d, \tprogress: %d%%%s") % (COLOR_BLUE, elapsed_time, 100, COLOR_BLACK),
            sys.stdout.flush()
        print ""
        (stdout, stderr) = p.communicate()
        for line in stdout.split('\n'):
            print line
        if p.returncode != 0:
            print "Failed!!!"
        else:
            print "Done successfully!!!"