path: root/liveusb/creator.py

# -*- coding: utf-8 -*-
#
# Copyright © 2008  Red Hat, Inc. All rights reserved.
# Copyright © 2009  Mandriva. All rights reserved.
#
# This copyrighted material is made available to anyone wishing to use, modify,
# copy, or redistribute it subject to the terms and conditions of the GNU
# General Public License v.2.  This program is distributed in the hope that it
# will be useful, but WITHOUT ANY WARRANTY expressed or implied, including the
# implied warranties 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., 51 Franklin Street, Fifth
# Floor, Boston, MA 02110-1301, USA. Any Red Hat trademarks that are
# incorporated in the source code or documentation are not subject to the GNU
# General Public License and may only be used or replicated with the express
# permission of Red Hat, Inc.
#
# Author(s): Luke Macken <lmacken@redhat.com>
#            Aurelien Lefebvre <alefebvre@mandriva.com>
#	     Luis Medinas <luis.medinas@caixamagica.pt>

import subprocess
import logging
import os
import re
import signal

from datetime import datetime
from stat import ST_SIZE
from threading import Thread
from time import sleep

from liveusb import _

class LiveUSBError(Exception):
    _message = ""
    def __init__(self, message): self._message = message
    def _get_message(self): return self._message
    def _set_message(self, message): self._message = message
    message = property(_get_message, _set_message)

class LiveUSBCreator(object):
    """ An OS-independent parent class for Live USB Creators """

    iso = None
    drives = {}
    pids = []
    isosize = 0
    log = None
    drive = None

    def __init__(self, opts):
        self.opts = opts
        self._setup_logger()

    def _setup_logger(self):
        self.log = logging.getLogger(__name__)
        level = logging.INFO
        if self.opts.verbose:
            level = logging.DEBUG
        self.log.setLevel(level)
        self.handler = logging.StreamHandler()
        self.handler.setLevel(level)
        formatter = logging.Formatter("[%(module)s:%(lineno)s] %(message)s")
        self.handler.setFormatter(formatter)
        self.log.addHandler(self.handler)

    def detect_removable_drives(self):
        """ This method should populate self.drives with removable devices """
        raise NotImplementedError

    def terminate(self):
        """ Terminate any subprocesses that we have spawned """
        raise NotImplementedError

    def set_iso(self, iso):
        """ Select the given ISO """
        self.iso = self._to_unicode(iso)
        self.isosize = os.stat(self.iso)[ST_SIZE]

    def _to_unicode(self, obj, encoding='utf-8'):
        if hasattr(obj, 'toUtf8'): # PyQt4.QtCore.QString
            obj = str(obj.toUtf8())
        #if isinstance(obj, basestring):
        #    if not isinstance(obj, unicode):
        #        obj = unicode(obj, encoding, 'replace')
        return obj

    def _test_hybrid_1(self, iso):
        hybrid = False
        for i in range(0,512):
            if not iso.read(1) == '\x00':
                hybrid = True
            if hybrid:
                break
        return hybrid
    
    def _test_hybrid_2(self, iso):
        hybrid = False
        iso.seek(0x1fe)
        if iso.read(1) == '\x55':
            iso.seek(0x1ff)
            if iso.read(1) == '\xaa':
                hybrid = True
        return hybrid
    
    def _test_hybrid_3(self, iso):
        hybrid = True
        iso.seek(0x200)
        for i in range(0x200,0x8000):
            if iso.read(1) != '\x00':
                hybrid = False
                break
        return hybrid
    
    def is_hybrid(self, iso):
        isofile = open(iso, "rb")
        hybrid =  self._test_hybrid_1(isofile) and self._test_hybrid_2(isofile) and self._test_hybrid_3(isofile)
        isofile.close()
        return hybrid

class LinuxLiveUSBCreator(LiveUSBCreator):

    bus = None # the dbus.SystemBus
    hal = None # the org.freedesktop.Hal.Manager dbus.Interface

    def detect_removable_drives(self):
        """ Detect all removable USB storage devices using HAL via D-Bus """
        import dbus
        self.drives = {}
        self.bus = dbus.SystemBus()
        hal_obj = self.bus.get_object("org.freedesktop.Hal",
                                      "/org/freedesktop/Hal/Manager")
        self.hal = dbus.Interface(hal_obj, "org.freedesktop.Hal.Manager")

        devices = []
        devices = self.hal.FindDeviceByCapability("storage")

        for device in devices:
            try:
                dev = self._get_device(device)
                if dev.GetProperty("storage.bus") == "usb":
                    self._add_device(dev)
            except dbus.exceptions.DBusException:
                pass

        if not len(self.drives):
            raise LiveUSBError(_("Unable to find any USB drives"))

    def _add_device(self, dev, parent=None):
        model = None
        capacity = None
        try:
            model = dev.GetProperty('storage.model')
            capacity = str(dev.GetProperty('storage.removable.media_size'))
        except Exception:
            pass
        
        self.drives[dev.GetProperty('block.device')] = {
            'name'     : dev.GetProperty('block.device'),
            'model'    : model,
            'capacity' : capacity
        }

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

        isosize = float(self.isosize)
        pattern = "^([0-9]+) bytes .*"
        patternCompiled = re.compile(pattern)

        pidPattern = "^DDPID=([0-9]+).*"
        pidPatternCompiled = re.compile(pidPattern)

        drive_infos = self.drives[self.drive]
        if drive_infos.has_key("capacity") and drive_infos["capacity"]:
            self.log.debug("Iso size = %s" % str(isosize))
            self.log.debug("Device capacity = %s " % str(drive_infos['capacity']))
            if int(drive_infos['capacity']) < int(isosize):
                raise LiveUSBError(_("Selected iso is too large for the selected device"))

        dd = '/usr/sbin/liveusb-dd.sh "' + self.iso + '" "' + self.drive + '"'
        self.log.debug(dd)
        p = subprocess.Popen(dd, 
          shell=True, 
          stdout=subprocess.PIPE, 
          stderr=subprocess.STDOUT,
          env={"LC_ALL": "C"})

        ppid = p.pid
        ddpid = None

        self.pids.append(ppid)
        
        # Updating progress bar by parsing /bin/dd output
        while True:
            line = p.stdout.readline()
            if not line:
                break

            if not ddpid:
                pidm = re.findall(pidPatternCompiled, line)
                if len(pidm) > 0:
                    ddpid = int(pidm[0])
                    self.pids.append(ddpid)

            m = re.findall(patternCompiled, line)
            if len(m) > 0:
                current = float(m[0])
                progress = (current/isosize)*100
                progress = round(progress, 2)
                progress_thread.update(progress)

        self.pids.remove(ppid)
        self.pids.remove(ddpid)
        
    def terminate(self):
        for pid in self.pids:
            try:
                os.kill(pid, signal.SIGHUP)
                self.log.debug("Killed process %d" % pid)
            except OSError:
                pass

class WindowsLiveUSBCreator(LiveUSBCreator):

    def detect_removable_drives(self):
        import win32file, win32api, pywintypes, wmi
        self.drives = {}
                
        c = wmi.WMI()

        for physical_disk in c.Win32_DiskDrive ():
          if physical_disk.InterfaceType == "USB":
            if len(physical_disk.Capabilities):
              for objElem in physical_disk.Capabilities:
               if objElem == 7:
                self.drives[physical_disk.DeviceID] = {
			    'name'    : physical_disk.DeviceID,
			    'model'   : physical_disk.Model,
			    'device'  : physical_disk.DeviceID,
			    'capacity': physical_disk.Size
			    }
                break
        if not len(self.drives):
            raise LiveUSBError(_("There isn't any removable drive available"))

    def _get_device(self, drive):
        if not self.drives[drive]:
            return None
        return self.drives[drive]['device']
    
    def build_disk(self, progress_thread):

        isosize = float(self.isosize)
        pattern = "^([0-9,]+).*"
        patternCompiled = re.compile(pattern)
        
#        f = subprocess.Popen(['format', self.drive, '/q', '/y', '/fs:fat32'], shell=True)
#        self.log.debug("Formating %s" % self.drive)
#        f.wait()
#        self.log.debug("Done, exit code: %d" % f.wait())

        dd = os.path.join("tools", "dd.exe")
        dd += " bs=8M" + " if=\"%s\" of=%s --progress" % (self.iso, self._get_device(self.drive))
        self.log.debug(dd)
        p = subprocess.Popen(dd, 
                             shell=True, 
                             stdout=subprocess.PIPE, 
                             stderr=subprocess.STDOUT,
                             universal_newlines=True)

        ppid = p.pid

        self.pids.append(ppid)
        
        # Updating progress bar by parsing /bin/dd output
        while True:
            line = p.stdout.readline()
            if not line:
                break

            m = re.findall(patternCompiled, line)
            if len(m) > 0:
                current = float(re.sub(",", "", m[0]))
                progress = (current/isosize)*100
                progress = round(progress, 2)
                progress_thread.update(progress)

        self.pids.remove(ppid)
    
    def terminate(self):
        """ Terminate any subprocesses that we have spawned """
        import win32api, win32con, pywintypes
        for pid in self.pids:
            try:
                handle = win32api.OpenProcess(win32con.PROCESS_TERMINATE,
                                              False, pid)
                self.log.debug("Terminating process %s" % pid)
                win32api.TerminateProcess(handle, -2)
                win32api.CloseHandle(handle)
            except pywintypes.error:
                pass