#!/usr/bin/env python

## (c) 2009 Laurent Coustet
# Clarisys Informatique
#
# Script to resize a qcow2 qemu image

import struct
import decimal

qcow2header = struct.Struct(">IIQIIQIIQQIIQ") # Big Endian

def sizeof_fmt(num):
    bak = num
    for x in ['bytes','KB','MB','GB','TB']:
        if num < 1024.0:
            return "%d %3.1f%s" % (bak, num, x)
        num /= 1024.0

class QCow2Header(object):
	def __init__(self, *args):
		print "args: %s" % (args,)
		self.magic = args[0]
		self.version = args[1]
		self.backing_file_offset = args[2]
		self.backing_file_size = args[3]
		self.cluster_bits = args[4]
		self.size = args[5]
		self.crypt_method = args[6]
		self.l1_size = args[7]
		self.l1_table_offset = args[8]
		self.refcount_table_offset = args[9]
		self.refcount_table_clusters = args[10]
		self.nb_snapshots = args[11]
		self.snapshots_offset = args[12]
	def __str__(self):
		return "Qcow2Header %s version %s size: %s" % (self.magic, self.version, sizeof_fmt(self.size))
	def check(self):
		'Q' 'F' 'I' '\xfb'
		return bool((ord('Q') << 24 | ord('F') << 16 | (ord('I') << 8) | ord('\xfb') << 0) == self.magic)
	def resize(self, newsize):
		"""
		@params newsize: size in bytes K for Kilobytes M for Megabytes G for Gigabytes
		Exemple 200 = 200 Bytes
			200M = 200 * 1024 * 1024 Bytes
			2G = 2 * 1024 * 1024 * 1024 Bytes
			+50M = (currentSize + (50 * 1024 * 1024)) Bytes
		"""
		add = False
		newsize = newsize.upper()
		# Extract "+"
		if newsize[0] == "+":
			add = True
			newsize = newsize[1:]
		
		# Extract unit
		units = {"B": 1, "K": 1024,"M": 1024 ** 2, "G": 1024**3}
		unit = "B"
		for u in units.keys():
			if newsize[-1] == u:
				unit = u
				newsize = newsize[:-1]
				break
		newsize_int = int(newsize) * units[unit]
		if add:
			newsize_int += self.size
		print "New size: %s Old size: %s" % (sizeof_fmt(newsize_int), sizeof_fmt(self.size))
		self.size = newsize_int
		
		# Recalculate the number l1 size
		# LC: Uggly method
		size_mega = int(self.size / (1024 ** 2))
		self.l1_size = int(decimal.Decimal(decimal.Decimal(size_mega) / decimal.Decimal(512)).quantize(decimal.Decimal('1.'), rounding=decimal.ROUND_UP))
		print "l1_size: %s" % (self.l1_size,)
		
	def get(self):
		return qcow2header.pack(self.magic, self.version, self.backing_file_offset, self.backing_file_size, self.cluster_bits, self.size, self.crypt_method, self.l1_size, self.l1_table_offset, self.refcount_table_offset, self.refcount_table_clusters, self.nb_snapshots, self.snapshots_offset)
	

def help():
	print """QCow2 Image resizer (grow only)
	(c) 2009 Laurent Coustet <address@hidden>
	
	Usage: %s file new_size_in_bytes
	"""

if __name__ == "__main__":
	import sys
	if len(sys.argv) < 3:
		help()
		sys.exit(0)

	f = open(sys.argv[1], "rb")
	data = f.read(qcow2header.size)
	f.close()

	if len(data) < qcow2header.size:
		print "Invalid header. Not a qcow2 file. Must be %s long." % (qcow2header.size,)

	# Try to use unpack
	header = QCow2Header(*qcow2header.unpack(data))
	if not header.check():
		print "This is not a QCow2 image: magic not found."
		sys.exit(0)

	# Resize the header
	header.resize(sys.argv[2])

	# Rewrite the file
	f = open(sys.argv[1], "r+b")
	f.seek(0, 0)
	ret = f.write(header.get())
	f.close()