#!/usr/bin/python -ttu
# vim: ai ts=4 sts=4 et sw=4

#    Copyright (c) 2008 Intel Corporation
#    Copyright (c) 2009 Bert Freudenberg
#
#    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 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 sys
import os
import gettext
import re
import subprocess
import signal
import time
import select

_ = gettext.lgettext
PROG_NAME = "image-writer-mac"
COLOR_BLACK = "\033[00m"
COLOR_RED =   "\033[1;31m"
COLOR_BLUE =  "\033[1;34m"

def main():

    # Python version check
    if sys.hexversion < 0x2040000:
        print >> sys.stderr, _("Error: %s depends on a Python " \
                               "v2.4 or greater!") % (PROG_NAME)
        sys.exit(1)

    # Parameters check
    if len(sys.argv) != 2:
        display_help()
        sys.exit(1)

    # Check image validity
    image = os.path.realpath(os.path.abspath(sys.argv[1]))
    if not is_valid_image(image):
        sys.exit(1)

    # Get USB drive
    usbd=get_usb_disk()
    if len(usbd)==0:
        sys.exit(1)

    # Unmount USB drive
    if not umount_device(usbd):
        sys.exit(1)

    # Warn of pending USB content doom
    msg = _("Warning:  The USB drive (%s) will be completely erased!") % diskutil_name(usbd)
    if not user_confirm (msg, COLOR_RED):
        sys.exit(1)

    # Write image to USB drive
    if not write_image_to_disk (image, usbd):
        sys.exit(1)

    print _("You may now boot your mobile device with this USB drive\n")
    return 0

#-----------------------------------------------------------------------

def display_help():
    print _("\nUsage: %s <full-path-to-image>\n") % PROG_NAME

# ret True = running as root
def are_we_root():
    if os.getuid() == 0:
        return True
    return False

def is_valid_image(image):
    if not os.path.isfile(image):
        print _("The image path does not point to a file")
        return False;
    if not (os.path.getsize(image) > (180*1024*1024)):
        msg = _("The file specified does not appear to be " \
                "large enough to be a valid image.")
        if not user_confirm (msg):
            return False;
    return True

def user_confirm(msg, text_color=COLOR_BLACK,
                 confirm=_(" Do you want to continue anyway? (y/n) ")):
    print "%s%s%s" % (text_color, msg, COLOR_BLACK)
    confirm_msg = "%s%s%s" % (text_color, confirm, COLOR_BLACK)
    name=""
    while name!="n" and name!="y":
        name = raw_input(confirm_msg)
        name = name.lower()
    if name=="n":
        return False;
    return True

def get_usb_disk():
    usbd = ""
    udisks = get_current_udisks();
    if len(udisks)==0:
        print _("No USB drives detected.")
        print _("Please insert a USB drive large enough to store the " \
                "given image")
        return ""
    elif len(udisks)==1:
        usbd=udisks[0]
    elif len(udisks)>0:
        print _("\nMultiple USB drives discovered:")
        i=1
        for usbd in udisks:
            print "%s) %s" % (i, diskutil_name(usbd))
            i+=1
        sel_usb = ""
        while sel_usb not in range (1,len(udisks)+1):
            try:
                sel_usb = int(raw_input(" Select the USB " \
                                        "drive to use (1-%s): " % (len(udisks))))
            except ValueError:
                continue
        usbd = udisks[sel_usb-1]
    print "%sDrive selected: %s%s" % (COLOR_BLUE, diskutil_name(usbd), COLOR_BLACK)
    return usbd

def diskutil(*args):
    return subprocess.Popen(("diskutil",) + args, stdout=subprocess.PIPE).communicate()[0].splitlines()

def diskutil_list():
    return [d for d in diskutil("list") if d.startswith("/")]

def diskutil_info(disk):
    info = {}
    for line in diskutil("info", disk):
        try:
            (key, value) = line.split(":", 1)
            info[key.strip()] = value.strip()
        except:
            pass
    return info

diskutil_name_cache = {}
def diskutil_name(disk):
    if disk not in diskutil_name_cache:
        info = diskutil_info(disk)
        name = info['Volume Name']
        media = info['Device / Media Name']
        size = info['Total Size']
        if name.startswith("Not applicable"):
            name = ""
        if name == "":
            for line in diskutil("list", disk):
                try:
                    (partition, description) = line.split(":", 1)
                    if partition.strip() != '#':
                        if name != "":
                            name += ","
                        name += line[33:56].strip()
                except:
                    pass
        if name == "":
            name = "Untitled"
        size = " ".join(size.split()[:2])
        diskutil_name_cache[disk] = "%s [%s, %s, %s]" % (name, disk, media, size) 
    return diskutil_name_cache[disk]

def get_current_udisks():
    usb_devices = []
    print _("%sLooking for storage devices...%s") % (COLOR_BLUE, COLOR_BLACK)
    for disk in diskutil_list():
        info = diskutil_info(disk)
        protocol = info["Protocol"]
        name = diskutil_name(disk)
        print "%s: %s" % (protocol, name)
        if protocol == "USB":
            usb_devices.append(disk)
    return usb_devices

def umount_device(device):
    diskutil("unmountDisk", device)
    return True

def write_image_to_disk (image_filename, usb_disk):
    size = os.path.getsize(image_filename)
    print _("Source:      %s") % image_filename
    print _("Size:        %s MB") % (int(size/(1024*1024)))
    print _("Destination: %s") % diskutil_name(usb_disk)
    #(tbd, get usb capacity): print _("Capacity:    %s") % ("1GB")

    msg = _("Writing image ...")
    print "%s%s%s" % (COLOR_BLUE, msg, COLOR_BLACK) ,

    cmd = ["dd", "bs=262144", "if=%s" % image_filename, "of=%s" % usb_disk]
 
    p = subprocess.Popen(cmd,
                         stdout = subprocess.PIPE,
                         stderr = subprocess.PIPE,
                         close_fds = True)

    # show progress (percentage ticking)
    while p.poll() == None:
        time.sleep(1)
        os.kill(p.pid, signal.SIGINFO)
        recordsIn = p.stderr.readline().strip()
        recordsOut = p.stderr.readline().strip()
        bytesTransferred = p.stderr.readline().strip()
        try:
            perc = int(bytesTransferred.split()[0]) * 100 / size
            print "%s\b\b\b\b%2s%%%s" % (COLOR_BLUE, perc, COLOR_BLACK) ,
        except:
            print recordsIn
            print recordsOut
            print bytesTransferred
            pass
    if p.returncode == 0:
        print "%s\b\b\b\b\b100%%%s" % (COLOR_BLUE, COLOR_BLACK)
    else:
        print

    # show errors or output of command
    (sout,serr) = p.communicate()
    for line in serr.split('\n'):
        print line
    for line in sout.split('\n'):
        print line

    if p.returncode != 0:
        print _("Error:  The image could not be written to the USB drive")
        return False
    print _("The image was successfully written to the USB drive")
    return True


if __name__ == '__main__':
    sys.exit(main())