gpthybrid

#!/usr/bin/python -tt
# -*- coding: utf-8 -*-

from subprocess import *
from optparse import OptionParser

parser = OptionParser()
parser.add_option("-l", "--label", dest="l", default="", help="set the partlabel stem for rootfs partitions")
parser.add_option("-i", "--iso", dest="i", help="name of the iso to gpthybridise")
parser.add_option("-s", "--sector", dest="s", default=4, help="size of sectors to use in resulting partition table")
(options, args) = parser.parse_args()

iso_file = options.i
label = options.l
sectors = options.s

def modup(number,mod,target):
	if (number % mod != target):
		number = number + target - number % mod
	return number
def moddown(number,mod,target):
	if (number % mod != target):
		adjust = number % mod - target
		if (adjust > 0):
			number -= adjust
		else:
			number -= mod + adjust
	return number

buffout = "gpthybridising iso %s" % iso_file
if len(label):
	buffout = "%s with label %s" % (buffout, label)
print buffout

# extract start and end of space to partition from existing gpt
gdisks = Popen(["/sbin/gdisk","-l",iso_file],stdout=PIPE).communicate()[0]
last = 0
end = 0
i = 0
for old_part in gdisks.split("\n"):
	if old_part.startswith('   1'):
		# below we want last as "end" of space before 1st partiion
		last = int(old_part.split()[1])-1
		# min 63 and round up to %sectors==sectors-1 (so 63 means next starts on 64)
		if last<63:
			last=63
		last = modup(last,sectors,sectors-1)
		print "first partition set last to %i" % last
	if old_part.startswith('   '):
		# end should be the last used sector
		testend = int(old_part.split()[2])
		if (testend > end):
			end = modup(testend,sectors,sectors-1)
			print "got partition ending: %i" % end
		i += 1

# get the lba addresses of the files using osirrox
parts = list()
osirrox=list(["/usr/bin/osirrox", "-pkt_output", "on",
		 "-indev", iso_file,
		 '-logfile', '.', '-',])
		 #'-find', '/efi.img', '-exec', 'report_lba', '--'])

for file in args:
	osirrox.extend(['-find', file, '-exec', 'report_lba', '--'])

print "running: %s" % " ".join(osirrox)
for line in check_output(osirrox).split("\n"):
	if line[:22] == 'R:1: File data lba:  0':
		# extract the fields
		(t, s, b, f, n) = line.split(',')
		# osirrox talks in 2048b sectors so *4
		s = int(s.split()[0])*4
		b = int(b.split()[0])*4
		# drop the quotes around the name
		n = n.split()[0].split("'")[1]
		# rebuild list in post in order
		done = 0
		post = list()
		for part in (parts):
			if s < int(part.split()[0]):
				# we start before this part
				if done == 0:
					# add ourselves if not done
					post.append("%s %s %s" % (moddown(s,sectors,0), modup(s+b-1,sectors,3), n))
					done = 1
			post.append(part)
		if done == 0:
			# no parts for us to go before so add to the end
			post.append("%s %s %s" % (moddown(s,sectors,0), modup(s+b-1,sectors,3), n))
		parts = post

# gdisk commands to delete the existing partitions
cmddel = ""
# all > 1 need the number
while (i>1):
	cmddel += """d
%i
""" % i
	i -= 1
# the first partition needs no number
if i==1:
	cmddel += "d\n"

# gdisk commands to create the partitions
# gap and pnum just count
gap = 1
pnum = 1
# set sectors alignment (2048b iso sector size by default) for partition creation
cmdpart = """x
l
%i
m
""" % sectors
for part in parts:
	(s, e, n) = part.split()
	# default to ms basic partition and no custom label
	t = '0700'
	l=''
	if len(n) > 4 and ( n[-6:] == '.amd64' or n[-4:] == '.686' or n[-2:] == '.2' ):
		t = '8300'
		# use label.extension or just the filename as the partlabel
		if len(label)==0:
			l = '%s' % n
		else:
			l = '%s.%s' % (label, n[n.rfind('.')+1:])
	elif len(n) > 6 and n[-7:] == 'efi.img':
		t = 'ef00'
	elif len(n) >5 and n[-5:] == '.ef02':
		t = 'ef02'
	# if the last partition ended more then sectors less then this starts
	if int(s)-last >= sectors:
		cmdpart += """n
%i
%i
%i
0700
c
""" % (pnum, modup(last+1,sectors,0), moddown(int(s)-1,sectors,3))
		if pnum > 1:
			cmdpart += "%s\n" % str(pnum)
		cmdpart += "Gap%i\n" % (gap)
		pnum += 1
		gap += 1
	# create the partition itself
	cmdpart += """n
%i
%s
%s
%s
""" % (pnum, s, e, t)
	if len(l) > 0:
		cmdpart += """c
%i
%s
""" % (pnum, l)
	pnum += 1
	last = int(e)
	cmdpart += """p
"""

# create a gap from end of last partition and end of last gap if needed
if (end-last>sectors-1):
	cmdpart += """n
%i
%s

0700
c
%i
Gap%i
p
""" % (pnum,modup(last+1,sectors,0), pnum, gap)

# set back to normal 2048 sector alignment for new partitions
cmdpart += """x
l
2048
m
"""

# confirm writing the partition table
cmdend = """w
Y
"""

# run our commands through gdisk
cmd = "%s\n%s\n%s" % (cmddel, cmdpart, cmdend)
print "gdisk stdin: %s" % cmd.replace("\n", "\\n")
gdisks = Popen(["/sbin/gdisk",iso_file],stdin=PIPE).communicate(input=cmd)

# TODO read gdisk result and verify or throw big warning/error