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>
#               Yves Brungard

import logging
import os
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.log.info("File "+self.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
	iface = None # the org.freedesktop.Udisks.Manager dbus.Interface

	def detect_removable_drives(self):
		""" Detect all removable USB storage devices using UDisks 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")
		proxy = self.bus.get_object("org.freedesktop.UDisks", "/org/freedesktop/UDisks")
		self.iface = dbus.Interface(proxy, "org.freedesktop.UDisks")
		
		devs = self.iface.EnumerateDevices()
		
		for dev in devs:
			dev_obj = self.bus.get_object("org.freedesktop.UDisks", dev)
			dev = dbus.Interface(dev_obj, "org.freedesktop.DBus.Properties")
		
			if str(dev.Get('', 'DriveConnectionInterface')) == 'usb' and not str(dev.Get('', 'PartitionType')) and str(dev.Get('', 'DeviceIsMediaAvailable')) == '1':
				self._add_device(dev)

#		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 = str(dev.Get('', 'DriveModel'))
		 capacity = str(dev.Get('', 'DeviceSize'))
		 name= str(dev.Get('', 'DeviceFile'))
#		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
		 self.drives[name] = {
			'name'	 : name,
			'model'	: model,
			'capacity' : capacity
		}

	
	def build_disk(self, progress_thread):

		isosize = float(self.isosize)
		source = self.iso
		target = self.drive
		capacity=self.drives[self.drive]['capacity']
		if self.backup.get_active() :
			backup_dest=self.backup_select.get_label()
			self.log.info(_('Backup in:')+' '+backup_dest)
		self.log.info(_('Image: ')+source)
		self.log.info(_('Target Device: ')+target)
		b = os.path.getsize(source)
		if isosize > capacity :
			raise LiveUSBError(_("Selected iso is too large for the selected device"))
		else:
				if capacity> 1024*1024*1024*32 :
					raise LiveUSBError(_("Selected device is too big (>32Mb) to be a USB stick"))
				self.do_umount(target)
				# Writing step
				self.raw_write(source, target, b)
	   #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 = 'tools/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 raw_write(self, source, target, b, progress_thread):
		import io
		bs=4096*128
		try:
			ifc=io.open(source, "rb",1)
		except:
			raise LiveUSBError(_('Reading error.'))
		else:
			try:
				ofc= io.open(target, 'wb',0)
			except:
				 raise LiveUSBError(_('You have not the rights for writing on the device'))
			else:
				self.log.debug(_('Executing copy from ')+source+_(' to ')+target)
				steps=range(0, b+1, b/100)
				indice=1
				written=0
				ncuts=b/bs
				for i in xrange(0,ncuts):
					try:
						buf=ifc.read(bs)
					except:
						raise LiveUSBError(_("Reading error."))
					try:
						ofc.write(buf)
					except:
						raise LiveUSBError(_("Writing error."))
					written= written+bs
					if written > steps[indice]:
						if indice%5==0:
							self.log.info(_('Wrote: ')+str(indice)+'% '+str(written)+_(' bytes'))
						progress_thread.update(indice)
						indice +=1
						try:
							os.fsync(ofc)
						except:
						   raise LiveUSBError(_("Writing error."))
				progress_thread.update(100)
				self.log.debug(_('Image ')+source.split('/')[-1]+_(' successfully written to')+target)
				self.log.debug(_('Bytes written: ')+str(written))
				try:
					ofc.close()
				except:
					raise LiveUSBError(_("Writing error."))
			ifc.close()
		
	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
		import 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)
		source = self.iso
		target = self.drive
		capacity=eval(self.drives[self.drive]['capacity'])
		self.log.info(_('Image: ')+source)
		self.log.info(_('Target Device: ')+self.drive)
		b = os.path.getsize(source)
		if isosize > capacity :
			raise LiveUSBError(_("Selected iso is too large for the selected device"))
		else:
				if capacity> 1024*1024*1024*32 :
					raise LiveUSBError(_("Selected device is too big (>32Mb) to be a USB stick"))
				# Writing step
				self.raw_write(source, target, b, progress_thread)

	def raw_write(self, source, target, b, progress_thread):
		import io
		import win32con
		import win32file, win32event
		import winioctlcon
		bs=4096*128l
		try:
			ifc=io.open(source, "rb",1)
		except:
			raise LiveUSBError(_('Reading error.'))
		else:
			try:
				ofc= win32file.CreateFile(target,
                         win32con.GENERIC_WRITE|win32con.GENERIC_WRITE,
                         win32con.FILE_SHARE_READ | win32con.FILE_SHARE_WRITE,
                         None,
                         win32file.OPEN_EXISTING,
                         0,
                         None)
				win32file.DeviceIoControl(ofc, winioctlcon.FSCTL_LOCK_VOLUME,  None, 0,None)
				win32file.DeviceIoControl(ofc, winioctlcon.FSCTL_DISMOUNT_VOLUME,  None, 0,None )
			except:
				 raise LiveUSBError(_('You have not the rights for writing on the device'))
			else:
				steps=range(0, b+1, b/100)
				steps.append(b)
				indice=1
				written=0
				ncuts=b/bs
				writeOvlap = win32file.OVERLAPPED()
				writeOvlap.hEvent = win32event.CreateEvent(None, 0, 0, None)
				writeOvlap.Offset = 0l
				while ncuts <= 100:
				    bs=bs/2
				    ncuts=b/bs
				for i in xrange(0,ncuts+1):
					try:
						buf=ifc.read(bs)
					except:
						raise LiveUSBError(_("Reading error."))
					try:
        					win32file.WriteFile(ofc, buf, writeOvlap)
        					win32event.WaitForSingleObject(writeOvlap.hEvent, win32event.INFINITE)
        					writeOvlap.Offset = writeOvlap.Offset + bs
					except:
						raise LiveUSBError(_("Writing error or access denied ")+str(written))
					written= written+bs
					if written > steps[indice]:
					   progress_thread.update(indice)
					   indice =indice+1        
				try:
					ofc.close()
				except:
					raise LiveUSBError(_("Writing error"))
				ifc.close()
	
	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