View & Modify GPT tables

gpt.py

#!/usr/bin/env python
'''GUID Partition Table Handling routines

    GUID Partition Table is a new partitioning standard that was introduced
    by Intel with the proposal of the EFI. GPT is designed to overcome the
    limitations of the ancient MBR that day by day are facing the daylight.

    This module was designed to provide an API to direct access GPT structure.
    It can be used to create partitioning programmas or disaster recovery
    ones.

    Copyright (C) 2009 <The Still Anonymous Greek Hacker Space>

    Licence: GPLv3
    --------------------------------------------------------------------
    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, either version 3 of the License, or
    (at your option) any later version.

    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, see <http://www.gnu.org/licenses/>.
'''

import sys, math, zlib, ctypes, random, getopt
from struct import *
from array import *


# GPT Partition Types Description
partition_type_desc =	{
	'00000000-0000-0000-0000-000000000000':'Unused Entry',
	'024DEE41-33E7-11D3-9D69-0008C781F39F':'MBR partition scheme',
	'C12A7328-F81F-11D2-BA4B-00A0C93EC93B':'EFI System Partition',
	'21686148-6449-6E6F-744E-656564454649':'BIOS Boot Partition',
	'E3C9E316-0B5C-4DB8-817D-F92DF00215AE':'Microsoft Reserved Partition',
	'EBD0A0A2-B9E5-4433-87C0-68B6B72699C7':'Basic Data Partition (Windows/Linux)',
	'5808C8AA-7E8F-42E0-85D2-E1E90434CFB3':'Logical Disk Manager metadata partition (Windows)',
	'AF9B60A0-1431-4F62-BC68-3311714A69AD':'Logical Disk Manager data partition (Windows)',
	'75894C1E-3AEB-11D3-B7C1-7B03A0000000':'Data partition (HP/UX)',
	'E2A1E728-32E3-11D6-A682-7B03A0000000':'Service Partition (HP/UX)',
	'A19D880F-05FC-4D3B-A006-743F0F84911E':'Raid Partition (Linux)',
	'0657FD6D-A4AB-43C4-84E5-0933C84B4F4F':'Swap partition (Linux)',
	'E6D6D379-F507-44C2-A23C-238F2A3DF928':'Logical Volume Manager (LVM) partition (Linux)',
	'8DA63339-0007-60C0-C436-083AC8230908':'Reserved (Linux)',
	'83BD6B9D-7F41-11DC-BE0B-001560B84F0F':'Boot Partition (FreeBSD)',
	'516E7CB4-6ECF-11D6-8FF8-00022D09712B':'Data partition (FreeBSD)',
	'516E7CB5-6ECF-11D6-8FF8-00022D09712B':'Swap partition (FreeBSD)',
	'516E7CB6-6ECF-11D6-8FF8-00022D09712B':'Unix File System (UFS) partition (FreeBSD)',
	'516E7CB8-6ECF-11D6-8FF8-00022D09712B':'Vinum volume manager partition (FreeBSD)',
	'516E7CBA-6ECF-11D6-8FF8-00022D09712B':'ZFS partition (FreeBSD)',
	'48465300-0000-11AA-AA11-00306543ECAC':'Hierarchical File System (HFS+) partition (OSX)',
	'55465300-0000-11AA-AA11-00306543ECAC':'Apple UFS (OSX)',
	'6A898CC3-1DD2-11B2-99A6-080020736631':'ZFS (OSX)',
	'52414944-0000-11AA-AA11-00306543ECAC':'Apple RAID partition (OSX)',
	'52414944-5F4F-11AA-AA11-00306543ECAC':'Apple RAID partition, offline (OSX)',
	'426F6F74-0000-11AA-AA11-00306543ECAC':'Apple Boot partition (OSX)',
	'4C616265-6C00-11AA-AA11-00306543ECAC':'Apple Label (OSX)',
	'5265636F-7665-11AA-AA11-00306543ECAC':'Apple TV Recovery partition (OSX)',
	'6A82CB45-1DD2-11B2-99A6-080020736631':'Boot partition (Solaris)',
	'6A85CF4D-1DD2-11B2-99A6-080020736631':'Root partition (Solaris)',
	'6A87C46F-1DD2-11B2-99A6-080020736631':'Swap partition (Solaris)',
	'6A8B642B-1DD2-11B2-99A6-080020736631':'Backup partition (Solaris)',
	'6A898CC3-1DD2-11B2-99A6-080020736631':'/usr partition (Solaris)',
	'6A8EF2E9-1DD2-11B2-99A6-080020736631':'/var partition (Solaris)',
	'6A90BA39-1DD2-11B2-99A6-080020736631':'/home partition (Solaris)',
	'6A9283A5-1DD2-11B2-99A6-080020736631':'EFI_ALTSCTR',
	'6A945A3B-1DD2-11B2-99A6-080020736631':'Reserved Partition 1 (Solaris)',
	'6A9630D1-1DD2-11B2-99A6-080020736631':'Reserved Partition 2 (Solaris)',
	'6A980767-1DD2-11B2-99A6-080020736631':'Reserved Partition 3 (Solaris)',
	'6A96237F-1DD2-11B2-99A6-080020736631':'Reserved Partition 4 (Solaris)',
	'6A8D2AC7-1DD2-11B2-99A6-080020736631':'Reserved Partition 5 (Solaris)',
	'49F48D32-B10E-11DC-B99B-0019D1879648':'Swap partition (NetBSD)',
	'49F48D5A-B10E-11DC-B99B-0019D1879648':'FFS (NetBSD)',
	'49F48D82-B10E-11DC-B99B-0019D1879648':'LFS (NetBSD)',
	'49F48DAA-B10E-11DC-B99B-0019D1879648':'Raid Partition (NetBSD)',
	'2DB519C4-B10F-11DC-B99B-0019D1879648':'Concatenated Partition (NetBSD)',
	'2DB519EC-B10F-11DC-B99B-0019D1879648':'Encrypted Partition (NetBSD)'
	}

def guid_raw_to_txt(raw):
	''' Convert the GPT GUID from a 16 bytes array to a human readable string. 
		The GPT GUID is repsented in XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX and
		the first 8 bytes are stored in Little Endian while the last 8 are stored
		in big endian.
		Returns a human readable string of the GUID'''
	le_chunks = unpack('<LHH', raw[0:8])
	be_chunks = unpack('>HHL', raw[8:16])
	return "%08X-%04X-%04X-" % le_chunks + "%04X-%04X%08X" % be_chunks

def guid_txt_to_raw(txt):
	'''Convert a text represented GUID to a raw 16 bytes array. UNIMPLEMENTED!'''
	pass

def guid_generate():
	'''Generate a new GUID and return a raw 16 byte array.'''
	raw_guid = ""
	for i in range(0,16):
		raw_guid += chr(int(random.uniform(1,256)))
	return raw_guid
	
def ptype_get_description(guid_txt):
	''' Get the GUID of Partition Type (in text format) and return a description string of this partition type'''
	if partition_type_desc.has_key(guid_txt):
		return partition_type_desc[guid_txt]
	else:
		return 'Unknown'

class entry:
	'''Handler of a Partition Entry'''
	part_type = '\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
	part_guid = '\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
	start_lba = 0
	end_lba = 0
	part_name = "blabla"

	def size_in_lba(self):
		'''Calculate the size of partition in LBA units'''
		return self.end_lba - self.start_lba

	def size_in_byte(self):
		'''Calcualte the size of partition in Byte units'''
		return self.size_in_lba() * 512

	def parse_raw_data(self, raw):
		'''Parse raw data and extract information about the Partition Information '''
		self.part_type = raw[0:16]
		self.part_guid = raw[16:32]
		self.start_lba = unpack('<Q', raw[32:40])[0]
		self.end_lba = unpack('<Q', raw[40:48])[0]
		self.part_name = unpack('<72s', raw[56:128])[0]

	def part_type_text(self):
		'''Return the GUID of the partition type in text format'''
		return guid_raw_to_txt(self.part_type)

	def part_guid_text(self):
		'''Return the GUID of the partition in text format'''
		return guid_raw_to_txt(self.part_guid)

	def part_type_description(self):
		'''Return a description of this partition type if it is known'''
		return ptype_get_description(self.part_type_text())

class header:
	'''Handler of a Partition Header'''
	signature = "EFI PART"
	revision = (0, 0, 1, 0)
	header_size = 92
	header_checksum = 0
	reserved_1 = 0
	my_lba = 0
	alternative_lba = 0
	first_usable_lba = 0
	last_usable_lba = 0
	disk_guid= "\x00" * 16
	entry_array_lba = 2
	total_entries = 128
	entry_size = 128
	entry_array_checksum = int(0)
	reserved_2 = "\x00" * 420

	def usable_size_in_lba(self):
		'''Calculate the size of usable space for partitions in LBA units'''
		return self.last_usable_lba - self.first_usable_lba

	def usable_size_in_byte(self):
		'''Calcualte the size of usable space in Byte units'''
		return self.usable_size_in_lba() * 512

	def disk_guid_text(self):
		'''Return the GUID of the disk in text format'''
		return guid_raw_to_txt(self.disk_guid)

	def has_valid_signature(self):
		'''Check if the header has valid signature. Returns True if it is valid'''
		return self.signature == "EFI PART"

	def has_valid_signature_text(self):
		'''Check if the header has valid signature. Returns a string that says if string is valid'''
		if self.signature == "EFI PART":
			return "VALID!"
		else:
			return "This signature is NOT valid!"

	def has_valid_header_checksum(self):
		'''Check if the header has a valid checksum. Returns True if it is valid'''
		drender = self.__basic_render(0)
		crc32 = ctypes.c_uint32(0)
		crc32.value = zlib.crc32(drender)
		return crc32.value == self.header_checksum

	def has_valid_header_checksum_text(self):
		'''Check if the header has a valid checksum. Returns a string that says if string is valid'''
		if self.has_valid_header_checksum():
			return "VALID!"
		else:
			return "This checksum is NOT valid!"

	def has_valid_entries_checksum(self, entries_array):
		'''Check if the entries checksum is right.'''
		crc32 = ctypes.c_uint32(0)
		crc32.value = zlib.crc32(entries_array)
		if crc32.value == self.entry_array_checksum:
			return True

	def has_valid_entries_checksum_text(self, entries_array):
		'''Check if the entries checksum is right.'''
		if self.has_valid_entries_checksum(entries_array):
			return "VALID!"
		else:
			return "This checksum is NOT valid!"

	def parse_raw_data(self, raw):
		'''Parse raw data and extract information Partition Table Header '''
		self.signature = raw[0:8]
		self.revision = unpack('<4b', raw[8:12])
		self.header_size = unpack('<L', raw[12:16])[0]
		self.header_checksum = unpack('<L', raw[16:20])[0]
		self.reserved_1 = unpack('<L', raw[20:24])[0]
		self.my_lba = unpack('<Q', raw[24:32])[0]
		self.alternative_lba = unpack('<Q', raw[32:40])[0]
		self.first_usable_lba = unpack('<Q', raw[40:48])[0]
		self.last_usable_lba = unpack('<Q', raw[48:56])[0]
		self.disk_guid = raw[56:72]
		self.entry_array_lba = unpack('<Q', raw[72:80])[0]
		self.total_entries = unpack('<L', raw[80:84])[0]
		self.entry_size = unpack('<L', raw[84:88])[0]
		self.entry_array_checksum = unpack('<L', raw[88:92])[0]
		self.reserved_2 = raw[92:512]

	def __basic_render(self, checksum):
		bdata = pack("<8s4bLLLQQQQ16sQLLL", self.signature,
			self.revision[0],
			self.revision[1],
			self.revision[2],
			self.revision[3],
			self.header_size,
			checksum,
			self.reserved_1,
			self.my_lba,
			self.alternative_lba,
			self.first_usable_lba,
			self.last_usable_lba,
			self.disk_guid,
			self.entry_array_lba,
			self.total_entries,
			self.entry_size,
			self.entry_array_checksum
			)
		return bdata

	def render(self):
		'''Render the data of this header to a block ready to be stored in a LBA.
			
			The function returns a block of bytes 512 bytes long, which is ready
			to be stored at the hard disk. Even the checksum of the header is
			calculated.'''
		# Render without checksum and reserved 2
		bdata = self.__basic_render(0)
		# Calculate checksum
		crc32 = ctypes.c_uint32(0)
		crc32.value = zlib.crc32(bdata)
		# Store checksum
		fdata = bdata[0:16] + pack('<L', crc32.value) + bdata[20:92] + ("\x00" * 420)
		return fdata

	def create_entry_array_checksum(self, entry_array):
		crc32 = ctypes.c_uint32(0)
		crc32.value = zlib.crc32(entry_array)
		self.entry_array_checksum = crc32.value
		return self.entry_array_checksum
		

if __name__ == "__main__":

	def print_entry(e, count):
		'''Function to print entry information'''
		print "+-----------------| ENTRY %s |----------------------------------------------+" % count
		print "| Type       : " + "%-59s |" % (e.part_type_description())
		print "|              %-59s |" % (e.part_type_text())
		print "| Unique Id  : " + "%-59s |" % (e.part_guid_text())
		print "| Size       : " + "%-59s |" % (str(e.size_in_byte()/1024) + " KB")
		print "| Start LBA  : " + "%-59s |" % (str(e.start_lba))
		print "| End LBA    : " + "%-59s |" % (str(e.end_lba))
		print "| Name       : " + "%-59s " % (e.part_name)

	def print_header(h, title, entries_array):
		'''Function to print GPT Header Information'''
		print "+-------------------| %-29s |----------------------+" % title
		print "| Signature          : %-10s %-40s |" % (h.signature, "(%s)" % (h.has_valid_signature_text()))
		print "| Revision           : %-51s |" % ("%d.%d.%d.%d" % h.revision)
		print "| Header Size        : %-51s |" % ("%d" % h.header_size)
		print "| Header CRC32       : %-10s %-40s |" % (("%08X" % h.header_checksum), "(%s)" % (h.has_valid_header_checksum_text()))
		print "| Reserved           : %-51s |" % ("%08X" % h.reserved_1)
		print "| My LBA             : %-51s |" % (h.my_lba)
		print "| Other LBA          : %-51s |" % (h.alternative_lba)
		print "| First usable LBA   : %-51s |" % (h.first_usable_lba)
		print "| Last usable LBA    : %-51s |" % (h.last_usable_lba)
		print "| Total usable space : %-51s |" % (h.usable_size_in_byte())
		print "| Disk GUID          : %-51s |" % (h.disk_guid_text())
		print "| Entries start LBA  : %-51s |" % (h.entry_array_lba)
		print "| Total entries      : %-51s |" % (h.total_entries)
		print "| Size of entry      : %-51s |" % (h.entry_size)
		print "| Entry array CRC32  : %-10s %-40s |" % (("%08X" % h.entry_array_checksum),  "(%s)" % (h.has_valid_entries_checksum_text(entries_array)))
		print "|                       Reserved Space                                     |"


	def dev_seek_lba(lba):
		if lba > dev_total_lba:
			raise "InternalError, reading data out of disk space!"
		dev_handle.seek(lba*512,0)

	def dev_read_lba(lba, count = 1):
		dev_seek_lba(lba)
		return dev_handle.read(512 * count) 

	def print_entries(start_lba, total_entries, stop_on_unused = True):
		e = entry()
		# Show one by one
		count = 0
		while(count < total_entries):
			if count % 4 == 0:
				edata = dev_read_lba(start_lba + (count / 4))
			e.parse_raw_data(edata[128* (count % 4):128* (count % 4) + 128])
			if e.part_type_text() == "00000000-0000-0000-0000-000000000000":
				return
			print_entry(e, count)
			count +=1

	def action_list():
		h = header()		

		# Check size
		if dev_total_lba < 2:
			print "Device is too small!"
			return 1

		# Open LBA 1
		hdata = dev_read_lba(1)
		h.parse_raw_data(hdata)
		if not h.has_valid_signature():
			print "ERROR: Cannot find a valid GPT signature, if you are sure"
			print " that there are valid GPT data then try to use -d (dump) option."
			return 2
		edata = dev_read_lba(h.entry_array_lba, (h.total_entries * h.entry_size)/512)

		# Show primary GPT
		print_header(h, "Primary Header (@ 1 LBA))", edata)
		print_entries(h.entry_array_lba, h.total_entries, True)
		print "+--------------------------------------------------------------------------+"

		if exclude_secondary:
			return 0

		# Open the other GPT
		print
		print "                     .... partion data ..."
		print

		hdata = dev_read_lba(h.alternative_lba)
		h.parse_raw_data(hdata)
		if not h.has_valid_signature():
			print "WARNING: Although the Primary GPT was valid, the Backup GPT"
			print " is corrupted. Primary GPT is enough but a corrupted backup"
			print " GPT means that a general corruption may have occured on the device."
			return 2
		edata = dev_read_lba(h.entry_array_lba, (h.total_entries * h.entry_size)/512)

		# Show backup entries
		print_entries(h.entry_array_lba, h.total_entries, True)
		print_header(h, "Backup Entry (@ %s LBA))" % h.my_lba, edata)
		print "+--------------------------------------------------------------------------+"
		return 0

	def action_dump():
		h = header()		

		# Check size
		if dev_total_lba < 2:
			print "Device is too small!"
			return 1

		# Open Primary GPT at LBA 1
		hdata = dev_read_lba(1)
		h.parse_raw_data(hdata)
		edata = dev_read_lba(2,32)

		# Show primary GPT
		print_header(h, "Primary Header (@ 1 LBA))", edata)
		print_entries(2, 128, True)
		print "+--------------------------------------------------------------------------+"

		if exclude_secondary:
			return 0

		print
		print "                     .... partion data ..."
		print

		# Open the other GPT
		hdata = dev_read_lba(dev_total_lba - 1)
		h.parse_raw_data(hdata)
		edata = dev_read_lba(dev_total_lba - 33, 32)

		# Show backup entries
		print_entries(dev_total_lba - 33, 128, True)
		print_header(h, "Backup Entry (@ %s LBA))" % (dev_total_lba - 1), edata)
		print "+--------------------------------------------------------------------------+"
		return 0

	def action_fixheader():
		dh = header()
		h = header()
		dsk_guid = guid_generate()

		# Check size
		if dev_total_lba < 2:
			print "Device is too small!"
			return 1

		# Open entries
		edata = dev_read_lba(2,32)

		# Create primary header
		dh.disk_guid = dsk_guid
		dh.my_lba = 1
		dh.alternative_lba = dev_total_lba - 1
		dh.first_usable_lba = first_usable_lba
		dh.last_usable_lba = dev_total_lba - first_usable_lba
		dh.entry_array_checksum = h.create_entry_array_checksum(edata)
		hdata = dh.render()
		h.parse_raw_data(hdata)

		# Show header
		print_header(h, "Primary Header", edata)
		print "Header Size: %s bytes" % len(hdata)
		while(1):
			a = raw_input("This header will be placed to LBA 1 as a Primary GPT Header. Are you sure? (Y/n)")
			print a
			if a in ("Y", "y"):
				dev_wh = open(device, "rb+")
				dev_wh.seek(512*1,0)
				dev_wh.write(hdata)
				dev_wh.close()
				print "Done!"
				break 
			if a in ("N", "n"):
				break

		if exclude_secondary:
			return 0

		# Open backup entries
		edata = dev_read_lba(dev_total_lba - 33 ,32)

		# Create primary header
		dh = header()
		h = header()
		dh.disk_guid = dsk_guid
		dh.my_lba = dev_total_lba -1
		dh.alternative_lba = 1
		dh.first_usable_lba = first_usable_lba
		dh.last_usable_lba = dev_total_lba - first_usable_lba
		dh.entry_array_checksum = h.create_entry_array_checksum(edata)
		dh.entry_array_lba = dev_total_lba - 33
		hdata = dh.render()
		h.parse_raw_data(hdata)

		# Show header
		print_header(h, "Backup Header", edata)
		print "Header Size: %s bytes" % len(hdata)
		while(1):
			a = raw_input("This header will be placed to LBA %s as a Backup GPT Header. Are you sure? (Y/n)" % (str(dev_total_lba - 1)))
			print a
			if a in ("Y", "y"):
				dev_wh = open(device, "rb+")
				dev_wh.seek(-512, 2)
				dev_wh.write(hdata)
				dev_wh.close()
				print "Done!"
				break
			if a in ("N", "n"):
				break

		return 0

	def print_usage(msg):
		'''Print the usage of this program'''
		print "GUID Partition Table Mangler"
		print
		print " %s [-ldfh] <device>" % sys.argv[0]
		print
		print " Actions "
		print "  -l  List the partiitions of the device (Default action)"
		print "  -d  Dump all the sections that GPT info is even if it is"
		print "      not a valid GPT. This option is helpfull in case of"
		print "      disaster recovery"
		print "  -f  Restore GPT header. If you delete by mistake"
		print "      the GPT header using a partiotioning program,"
		print "      this option will try to restore ONLY the header"
		print "      of the GPT leaving untouch the GPT Partition"
		print "      Entries and the partitions itself."
		print "      ATTENTION: This options makes changes on the disk"
		print "      in a non-standard way, and it may destroy all your data!"
		print "  -h  Displays this message."
		print
		print "  --exclude-secondary"
		print "      Exclude the secondary (backup) GPT Table from this action."
		print "  --first-usable-lba=[lba]"
		print "      GPT Header holds the position of the first usable LBA. The"
		print "      EFI Specification says that this value cannot be smaller of 34"
		print "      but it gives freedom to use larger values. Microsoft implemented"
		print "      its disk utility to use a value of 34 while Apple uses a value"
		print "      of 40. The default here is 34."
		print
		if len(msg) > 0:
			print msg

	# Parse arguements and options
	optlist, args = getopt.getopt(sys.argv[1:], "ldfh", ["help", "exclude-secondary", "first-usable-lba="])
	# Check for argument
	if len(args) != 1:
		print_usage("ERROR: You must specify the device!")
		sys.exit(1)

	# Check for options
	mode = "list"
	exclude_secondary = False
	device = args[0]
	first_usable_lba = 34
	for o, a in optlist:
		if o in ("-h", "--help"):
			print_usage()
			sys.exit(0)
		elif o == "-l":
			mode = "list"
		elif o == "-d":
			mode = "dump"
		elif o == "-f":
			mode = "fixheader"
		elif o in "--exclude-secondary":
			exclude_secondary = True
		elif o == "--first-usable-lba":
			if int(a) < 34:
				print_usage("ERROR: The first usable LBA must be at least 34!")
				sys.exit(1)
			first_usable_lba = int(a)
		else:
			print_usage("ERROR: Unknown option \"%s\"" % o)
			sys.exit(1)

	# Open device
	dev_handle = open(device, "rb")
	dev_handle.seek(0,2)
	dev_size = dev_handle.tell()
	dev_total_lba = int(math.floor(float(dev_size) / float(512)))
	dev_handle.seek(0,0)

	# Do the action
	if mode == "list":
		ret = action_list()
	elif mode == "dump":
		ret = action_dump()
	elif mode == "fixheader":
		ret = action_fixheader()
	sys.exit(ret)