Das deutsche Python-Forum

Code: Alles auswählen

"""
	mbr switch v 0.01
	Switch any_to_first partition entry in mbr
"""
import struct
import sys
import win32api
import os
import win32com.client
import subprocess
import time

"""
    give user some information and ask to proceed
"""
print "WARNING!! Connect only USB-Stick to switch MBR !!"
# 1. Win32_DiskDrive
strComputer = "."
objWMIService = win32com.client.Dispatch("WbemScripting.SWbemLocator")
objSWbemServices = objWMIService.ConnectServer(strComputer,"root\cimv2")
colItems = objSWbemServices.ExecQuery("SELECT * FROM Win32_DiskDrive WHERE InterfaceType = \"USB\"")
DiskDrive_DeviceID = colItems[0].DeviceID
DiskDrive_Caption = colItems[0].Caption

print DiskDrive_DeviceID, " - ", DiskDrive_Caption

strComputer = "."
objWMIService = win32com.client.Dispatch("WbemScripting.SWbemLocator")
objSWbemServices = objWMIService.ConnectServer(strComputer,"root\cimv2")
colItems = objSWbemServices.ExecQuery("SELECT * FROM Win32_LogicalDisk WHERE DriveType = 2")
for objItem in colItems:
    if objItem.Caption != None:

		VolmeCptn = objItem.Caption
    if objItem.VolumeName != None:

		VolmeNme = objItem.VolumeName
print "Name: %s on %s " % (VolmeNme, VolmeCptn)


abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
gof = raw_input('Proceed with switching to MBR? [y] Insert number of Stick? [c] Abort {n] : ')
if gof != 'y':
	if gof == 'c':
		ins = raw_input("Insert number \\.\PhysicalDrive: ")
		inss = str(ins)
		mbr_fle = "\\\\.\\PhysicalDrive" + inss
	else:
		sys.exit(0)                    			# stopped


# --filenames--------------------------------------------------------------------------

mbr_new = VolmeNme
mbr_flew = "mbr-" + mbr_new + ".bin"
mbr_fle = DiskDrive_DeviceID
mbr_dd = "mbr_wrt.bat"
mbr_f = mbr_flew

# --read/write data--------------------------------------------------------------------
with open(mbr_fle, 'rb') as f:    			# read mbr in binary mode
	mbr = f.read(512)                 		# 512 if mbr
# --class mbr--------------------------------------------------------------------------
sgntr = mbr[510:]
if struct.unpack('<BB', sgntr)[0] != 85:    # check if: 55AA  or   AA55
	print 'I go, bye'
	sys.exit(0)                    	# no mbr --> nothing to do

prtswt = 0
strt = 446
end = 462
boot_strct = mbr[0:446]						# read to fill data for write
part_tble = []								# part. entry go in
part_sktr = []								# start sector of every part.
bootI = []									# maybe important later version

x = "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
for prttn in range(4):						# read part entry in list (Array)
	bt = struct.unpack('<B', mbr[strt:strt + 1])[0]
	if bt == 128:
		bootI.append("+")                   # Check if bootable
	else:
		bootI.append("-")                   # or not
	part_tble.append(mbr[strt:end]) 	    # read Partitiontable
	part_sktr.append(mbr[strt + 8:strt + 12]) # read startsectors of partitions
	strt = end                              # change start and end number
	end = end + 16
	if part_tble[prttn] != x:               # drop empty partition values
		prtswt += 1                         # increase for next loop

# -------------------------------------------------------------------------------------
f.close()
# --switch data------------------------------------------------------------------------
if prtswt < 2:
	sys.exit(0)                             # only 1 Partition --> nothing to switch

elif prtswt < 3:
	ins = 1                                 # 2 Partitions     --> switch 1, 2 => 2, 1

elif prtswt > 2:
	n = 0
	print 'Num of Partititions:', prtswt    # >2 Partitions     --> ask than switch
	while n <= 3:
		sektors = struct.unpack("<HH", part_sktr[n])[0]
		boot = bootI[n]
		nm = n + 1
		if sektors != 0:
			print "Partition: '%d'   StartSektor:  '%d' Boot: %s " % (nm, sektors, boot)
		n += 1
	ins = input("which partiton to switch?: ")
	ins -= 1                                # list goes 0 - 3, so reduce ins

# --execute switching------------------------------------------------------------------
part_tble.insert(0, part_tble[ins])        # [0,ins, 2, 3]     --> [ins, 0, 1, 2, 3]  -->  [ins, 1, 2, 3, 4]
del part_tble[ins + 1]                     # [ins, 1,"", 3, 4] --> [ins, 0, 1, 2, 3]  -->  [0, 1, 2, 3]
# --writing----------------------------------------------------------------------------
data = boot_strct + part_tble[0] + part_tble[1] + part_tble[2] + part_tble[3] + sgntr
ans = raw_input('Successfully switched. Write to MBR? [y/n] : ')
if ans == 'y':
	with open(mbr_flew, 'wb') as wrf:
		wrf.write(data)
		wrf.close()
# --write batch for dd.exe-------------------------------------------------------------

par1 = "dd.exe if=" + mbr_f + " of=\\\\?\Device\Harddisk" + \
			DiskDrive_DeviceID[-1] + "\Partition0 bs=512 count=1\n"
with open(mbr_dd, 'w') as cpy:
	seq = ["@echo off\n", "cd /d %~dp0\n", par1, "pause"]
	cpy.writelines( seq )
	cpy.close()

print mbr_dd, " is ready to start dd.exe!"
time.sleep(3)
# --ToDo:  run dd.exe with parameters: subprocess---------------------------------------
# --compile with pyinstaller------------------------------------------------------------
#C:\>pyinstaller -F swtchparttions_usb.py

Code: Alles auswählen

    mbr_new = VolmeNme
    mbr_flew = "mbr-" + mbr_new + ".bin"
    mbr_fle = DiskDrive_DeviceID
    mbr_dd = "mbr_wrt.bat"
    mbr_f = mbr_flew

Code: Alles auswählen

"""
    mbr switch v 0.01
    Switch any_to_first partition entry in mbr
"""
import struct
import sys

import win32com.client

MBR_SIZE = 512
MBR_PARTITION_TABLE_OFFSET = 446
MBR_PARTITION_ENTRY_SIZE = 16
MBR_SIGNATURE = '\x55\xaa'

BATCH_FILENAME = 'mbr_wrt.bat'
# 
# TODO Really just the last character of `drive_device_id`?
#   That looks suspicious.
# 
BATCH_TEMPLATE = '''\
@echo off
cd /d %~dp0
dd.exe if={mbr_filename} of=\\\\?\\Device\\Harddisk{drive_device_id[-1]}\\Partition0 bs={mbr_size} count=1
pause
'''


def main():
    print 'WARNING! Connect only USB-Stick to switch MBR!'

    wbem_locator = win32com.client.Dispatch('WbemScripting.SWbemLocator')
    wbem_server = wbem_locator.ConnectServer('.', r'root\cimv2')
    # 
    # TODO Replace ``*`` by actual column names we are interested in.
    # 
    rows = wbem_server.ExecQuery(
        'SELECT * FROM Win32_DiskDrive WHERE InterfaceType = "USB"'
    )
    # 
    # FIXME `IndexError` if there is no USB drive connected.
    # 
    drive_device_id = rows[0].DeviceID
    drive_caption = rows[0].Caption

    print drive_device_id, '-', drive_caption
    # 
    # FIXME This query needs a clause to make sure the result belongs to the
    #   first query's result.
    # 
    # TODO Replace ``*`` by actual column names we are interested in.
    # 
    rows = wbem_server.ExecQuery(
        'SELECT * FROM Win32_LogicalDisk WHERE DriveType = 2'
    )
    #
    # FIXME `NameError` after loop if `rows` is empty.
    # 
    for row in rows:
        if (row.Caption, row.VolumeName) != (None, None):
            volume_caption = row.Caption
            volume_name = row.VolumeName
            break
    print 'Name: %s on %s ' % (volume_name, volume_caption)

    answer = raw_input('Proceed with switching to MBR? [y] Abort [n]: ')
    if answer != 'y':
        sys.exit(0)

    with open(drive_device_id, 'rb') as raw_device:
        mbr = raw_device.read(MBR_SIZE)
    
    if mbr[-2:] != MBR_SIGNATURE:
        print 'I go, bye'
        sys.exit(0)

    empty_entry = '\x00' * MBR_PARTITION_ENTRY_SIZE
    partition_count = 0
    partition_table_entries = list()
    for i in xrange(4):
        offset = MBR_PARTITION_TABLE_OFFSET + i * MBR_PARTITION_ENTRY_SIZE
        entry = mbr[offset:offset + MBR_PARTITION_ENTRY_SIZE]
        partition_table_entries.append(entry)
        if entry != empty_entry:
            partition_count += 1

    if partition_count == 1:
        # 
        # Only 1 Partition --> nothing to switch.
        # 
        sys.exit(0)
    elif partition_count == 2:
        # 
        # FIXME Super uncool if second partiton entry is empty!
        # 
        switch_index = 1
    else:
        print 'Number of Partititions:', partition_count
        for i, entry in enumerate(partition_table_entries, 1):
            start_sector = struct.unpack('<HH', entry[8:12])[0]
            if start_sector != 0:
                is_bootable = (ord(entry[0]) & 127) != 0
                print "Partition: '%d'   StartSektor:  '%d' Boot: %s " % (
                    i, start_sector, '+' if is_bootable else '-'
                )
        #
        # FIXME Check user input!  Really!
        # 
        switch_index = int(raw_input('which partiton to switch? ')) - 1

    partition_table_entries.insert(
            0, partition_table_entries.pop(switch_index)
    )

    answer = raw_input('Write switch to MBR? [y/n]: ')
    if answer == 'y':
        mbr_filename = 'mbr-{}.bin'.format(volume_name)
        
        with open(mbr_filename, 'wb') as mbr_file:
            new_mbr = (
                mbr[:MBR_PARTITION_TABLE_OFFSET]
                + ''.join(partition_table_entries)
                + MBR_SIGNATURE
            )
            assert len(new_mbr) == MBR_SIZE
            mbr_file.write(new_mbr)

        with open(BATCH_FILENAME, 'w') as batch_file:
            batch_file.write(
                BATCH_TEMPLATE.format(
                    drive_device_id=drive_device_id,
                    mbr_filename=mbr_filename,
                    mbr_size=MBR_SIZE,
                )
            )

        print BATCH_FILENAME, 'is ready to start dd.exe!'


if __name__ == '__main__':
    main()

Code: Alles auswählen

with lock_volume(disk):