#!/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
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

if os.geteuid () != 0:
    errmsg("You must run mic-image-writer as root")
options = parse_options(sys.argv[1:])
if not options.console or options.gui:
    try:
        import pygtk
        pygtk.require('2.0')
        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.")

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")

class USBWriterWindow:
    def __init__(self, the_image_file):
        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.usb_disk = gtk.combo_box_new_text()
        self.usb_disk.set_size_request(80, -1)
        
        # Create a centering alignment object
        align = gtk.Alignment(0.50, 0.50, 0.9, 0.10)
        align.add(self.usb_disk)

        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.dbus_init()

        self.window.show_all()

    def dbus_init(self):
        """Initialize dbus"""
        self.bus = bus
        self.hal = hal
        self.usb_disks_num = 0
        self.detect_removable_drives()
        self.connect_hal_signal()
        self.refresh_usblist(self.drives)

    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 connect_hal_signal(self):
        # If we have access to HAL & DBus, intercept some useful signals
        self.hal.connect_to_signal('DeviceAdded',
                                            self.populate_devices)
        self.hal.connect_to_signal('DeviceRemoved',
                                            self.populate_devices)

    def populate_devices(self, *args, **kw):
        self.detect_removable_drives()
        self.refresh_usblist(self.drives)
        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.usb_disk.remove_text(0)

        for usb_disk in usblist:
            self.usb_disk.append_text(usb_disk)
            self.usb_disks.append(usb_disk)

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

    def detect_removable_drives(self):
        """ Detect all removable USB storage devices using HAL via D-Bus """
        self.drives = []

        devices = []
        if False:
            devices = self.hal.FindDeviceStringMatch('block.device',
                                                     self.opts.force)
        else:
            devices = self.hal.FindDeviceByCapability("storage")

        for device in devices:
            try:
                dev = self._get_device(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'))
                        self.drives.append(tmpdevice)
            except dbus.DBusException, e:
                continue

        if len(self.drives):
           self.refresh_usblist(self.drives)

    def _get_device(self, udi):
        """ Return a dbus Interface to a specific HAL device UDI """
        dev_obj = self.bus.get_object("org.freedesktop.Hal", udi)
        return dbus.Interface(dev_obj, "org.freedesktop.Hal.Device")

    def msgbox(self, msg):
        dlg = gtk.MessageDialog(None, 0, gtk.MESSAGE_WARNING, gtk.BUTTONS_OK, msg)
        dlg.set_title("Warning")
        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 _hal_unmount(self, udi):
        dev = self._get_device(udi)
        device = str(dev.GetProperty('block.device'))
        if dev.GetProperty('volume.is_mounted'):
            try:
                dev_obj = self.bus.get_object("org.freedesktop.Hal", udi)
                dev_obj.Unmount([""], dbus_interface="org.freedesktop.Hal.Device.Volume")
            except dbus.DBusException, msg:
                self.msgbox(device + ": " + str(msg))
                return -1
        return 0

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

        for partition in partitions:
            dev = self._get_device(partition)
            if dev.GetProperty("block.is_volume"):
                try:
                    device = str(dev.GetProperty('block.device'))
                    ret = self._hal_unmount(partition)
                    if ret != 0:
                        return -1
                except:
                    pass
            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 = self._get_device(child)
                        if dev.GetProperty("block.is_volume"):
                            device = str(dev.GetProperty('block.device'))
                            ret = self._hal_unmount(child)
                            if ret != 0:
                                return -1
        return 0

    def disable_some_widgets(self):
        self.image_file.set_sensitive(False)
        self.usb_disk.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.usb_disk.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.usb_disk.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]]
        ret = self.unmount_usb(self.usb_disks[usb_disk])
        if ret != 0:
            self.enable_some_widgets()
            return
        print argv
        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!!!")
        else:
            self.set_progress()
            self.msgbox("Done successfully!!!")
        self.enable_some_widgets()

    def read_image(self, widget):
        pass


def get_usb_list():
    drives = []

    devices = []
    devices = hal.FindDeviceByCapability("storage")
    for device in devices:
        try:
            dev_obj = 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.append({"device":tmpdevice, "vendor":tmpvendor, "product":tmpproduct})
        except dbus.DBusException, e:
            continue
    return drives

def hal_unmount(device):
    dev_obj = 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'))
            errmsg("Error: failed to unmount %s, %s" % (device, msg))
            
def unmount_usb(usbdisk):
    partitions = hal.FindDeviceStringMatch('block.device', usbdisk)

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

if __name__ == "__main__":
    if options.gui:
        USBWriterWindow(options.image_file)
        gtk.main()
    else:
        usb_list = get_usb_list()
        if not usb_list:
            print "Please insert your USB stick..."
        while not usb_list:
            time.sleep(2)
            usb_list = get_usb_list()
        index = 1
        print "Available usb disk:"
        for usbdisk in usb_list:
            print "\t[%d] %s" % (index, "%s: %s %s" % (usbdisk["device"], usbdisk["vendor"], usbdisk["product"]))
            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"
        unmount_usb(usb_list[choice-1]["device"])
        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" % usb_list[choice-1]["device"]]
        print "Source: %s" % options.image_file
        print "Target: %s" % usb_list[choice-1]["device"]
        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!!!"
