From b9231d4b3b5e5b808544157179e6e2ff101830eb Mon Sep 17 00:00:00 2001 From: Anne Nicolas Date: Wed, 4 May 2011 21:14:55 +0000 Subject: update files for usbdumper --- liveusb/creator.py | 309 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 309 insertions(+) create mode 100755 liveusb/creator.py (limited to 'liveusb/creator.py') diff --git a/liveusb/creator.py b/liveusb/creator.py new file mode 100755 index 0000000..0e19734 --- /dev/null +++ b/liveusb/creator.py @@ -0,0 +1,309 @@ +# -*- 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 +# Aurelien Lefebvre +# Luis Medinas + +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 -- cgit v1.2.1