// VMware VMDK Sparse Disk Utility
//
// Supported Compilers: Visual C++ .Net 2003, Visual C++ 2005, g++ 4.x
// Platforms: Windows 2000 or above, Linux with 2.6.x kernel or above
//
// Author : EaseWay (easeway@163.com)
// Date   : Oct. 2007
// Version: 1.0.1

// CHANGES
//
// 1.0.1 (Oct. 2007)
//    Add command 'convert' to update ddb.virtualHWVersion and ddb.adapterType
//
// 1.0.0 (Sep. 2007)
//    Initial version release

/* Copyright (c) 2007 EaseWay
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. The name of the author may not be used to endorse or promote products
 *    derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#ifndef _WIN32
#define _FILE_OFFSET_BITS	64
#endif

#include <iostream>
#include <string>
#include <map>
#include <list>
#include <vector>
#include <stdlib.h>
#include <time.h>
#include <fcntl.h>
#include <sys/stat.h>

#ifdef _WIN32
#include <io.h>
#define off_t		__int64
#define lseek64		_lseeki64
#define OPEN_FLAGS	(O_BINARY | O_RDWR | O_TRUNC | O_CREAT)
#define OPEN_RFLAGS	(O_BINARY | O_RDONLY)
#define OPEN_WFLAGS (O_BINARY | O_RDWR)
#define OPEN_MODE	(S_IREAD | S_IWRITE)
#else
#define tell(fd)	lseek(fd, 0, SEEK_CUR)
#define OPEN_FLAGS	(O_LARGEFILE | O_RDWR | O_TRUNC | O_CREAT)
#define OPEN_RFLAGS	(O_RDONLY)
#define OPEN_WFLAGS (O_RDWR)
#define OPEN_MODE	(S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH)
#endif

using namespace std;

// VMDK specific data types

#define VMDK_UINT	unsigned int
#define VMDK_CHAR	char
#define VMDK_BOOL	unsigned char
#define VMDK_SECT	unsigned long long

#ifdef _MSC_VER
#pragma pack(push, 1)
#endif

typedef struct
#ifndef _MSC_VER
__attribute__((packed))
#endif
_VMDK_SPEXTHEAD
{
	VMDK_UINT	magic;					// 'VMDK'
	VMDK_UINT	version;				// 1
	VMDK_UINT	flags;					// 3
	VMDK_SECT	capacity;				// capacity of extent in sectors
	VMDK_SECT	grain_size;				// size of grain in sectors
	VMDK_SECT	desc_off;				// descriptor offset in sectors
	VMDK_SECT	desc_size;				// descriptor size in sectors
	VMDK_UINT	gte_per_gt;				// number GT entries per GT
	VMDK_SECT	rgd_off;				// redundent GD offset in sectors
	VMDK_SECT	gd_off;					// GD offset in sectors
	VMDK_SECT	overhead;				// metadata size in sectors
	VMDK_BOOL	unclean;				// unclean shut down
	VMDK_CHAR	endline_chars[4];		// "\n \r\n"
} VMDK_SPEXTHEAD;

#ifdef _MSC_VER
#pragma pack(pop)
#endif

#define VMDK_SPEXTMAGIC		0x564d444b
#define VMDK_VERSION		1
#define VMDK_FLAGS			3
#define VMDK_ENDLINECHARS	"\n \r\n"
#define VMDK_CIDNOPARENT	0xffffffff

#define VMDK_GRAIN_SIZE	128
#define VMDK_GTE_NUM	512
#define VMDK_SEC_SIZE	512
#define VMDK_GT_SIZE	(VMDK_GTE_NUM * 4 / VMDK_SEC_SIZE)
#define VMDK_SPLIT_SIZE	0x3ff800

// Command Line Interface
static void usage(void)
{
	cout << "Usage: vmdkutil <command> [options] [infile] [outfile]" << endl
		 << endl
		 << "Options:" << endl
		 << "    -s, --size xxG|xxM        total disk size" << endl
		 << "    -p, --split               split into multiple files" << endl
		 << "    -m, --merge               merge into single file" << endl
		 << "    -c, --chains n            copy with parent files in n (n > 0 or n is 'all') levels" << endl
		 << endl
		 << "Commands:" << endl
		 << "    vmdkutil info infile" << endl
		 << endl
		 << "    vmdkutil convert file hwver=4|6 adapter=ide|buslogic|lsilogic" << endl
		 << endl
		 << "    vmdkutil create [options] outfile" << endl
		 << "        -s, --size required" << endl
		 << "        -p, --split optional" << endl
		 << "        -a, --adapter ide|buslogic|lsilogic" << endl
		 << "        -r, --hwver 5|6 hardware version" << endl
		 << endl
		 << "    vmdkutil clone [options] infile outfile" << endl
		 << "        -p, --split optional" << endl
		 << "        -m, --merge optional" << endl
		 << endl
		 << "    vmdkutil copy [options] infile outfile" << endl
		 << "        -p, --split optional" << endl
		 << "        -m, --merge optional" << endl
		 << "        -c, --chains optional" << endl
		 << endl;
	exit(1);
}

struct Opts
{
	unsigned long long disk_size;
	bool split;
	bool merge;
	int chains;
	string adapter;
	int hwver;
	
	static const int CHAIN_ALL = -1;
	
	Opts() 
	: disk_size(0), 
	  split(false), merge(false), chains(CHAIN_ALL),
	  adapter("lsilogic"), hwver(4)
	{ }
};

struct VMDKFileRef
{
	string access;
	VMDK_SECT capacity;
	string type;
	string filename;
	VMDK_SECT offset;
	
	VMDKFileRef() : access("RW"), capacity(0), type("SPARSE"), offset(0) { }
};

typedef list<VMDKFileRef> VMDKFileRefList;

struct VMDKAttrs
{
	VMDK_SECT metadata_size;
	VMDK_SECT grain_count;
	VMDK_SECT data_grain_count;
	VMDK_SECT gt_num;
	VMDK_SECT gd_size;
	VMDK_SECT desc_size;
	
	VMDKAttrs() { }
	VMDKAttrs(const Opts &opts, size_t desc_size = 0) { CalcAll(opts.disk_size / VMDK_SEC_SIZE, desc_size); }
	VMDKAttrs(const VMDK_SPEXTHEAD &head) { CalcAll(head.capacity, head.desc_size * VMDK_SEC_SIZE); }
	VMDKAttrs(VMDK_SECT capacity, size_t descsz) { CalcAll(capacity, descsz); }
	
	void CalcAll(VMDK_SECT capacity, size_t descsz)
	{
		data_grain_count = (capacity + VMDK_GRAIN_SIZE - 1) / VMDK_GRAIN_SIZE;
		gt_num  = (data_grain_count + VMDK_GTE_NUM - 1) / VMDK_GTE_NUM;
		gd_size = (gt_num * 4 + VMDK_SEC_SIZE - 1) / VMDK_SEC_SIZE;
		metadata_size = (VMDK_GT_SIZE * gt_num + gd_size) * 2;
		if (descsz)
		{
			desc_size = 20;//(descsz + VMDK_SEC_SIZE - 1) / VMDK_SEC_SIZE;
			metadata_size += desc_size;
		}
		else
			desc_size = 0;
	
		if (metadata_size & (VMDK_GRAIN_SIZE - 1)) 
		{
			metadata_size += VMDK_GRAIN_SIZE;
			metadata_size &= ~(VMDK_GRAIN_SIZE - 1);
		}
		grain_count = data_grain_count + metadata_size / VMDK_GRAIN_SIZE;
	}
	
	VMDK_SECT getFileSize() const
	{
		return grain_count * VMDK_GRAIN_SIZE;
	}
	
	VMDK_SECT getDataSize() const
	{
		return data_grain_count * VMDK_GRAIN_SIZE;
	}
	
	void ExportHead(VMDK_SPEXTHEAD &head) const
	{
		memset(&head, 0, sizeof(head));
		head.magic = VMDK_SPEXTMAGIC;
		head.version = VMDK_VERSION;
		head.flags = VMDK_FLAGS;
		head.capacity = getDataSize();
		head.grain_size = VMDK_GRAIN_SIZE;
		head.desc_off = desc_size ? 1 : 0;
		head.desc_size = desc_size;
		head.gte_per_gt = VMDK_GTE_NUM;
		head.rgd_off = head.desc_off ? head.desc_off + desc_size : 1;
		head.gd_off = head.rgd_off + (VMDK_GT_SIZE * gt_num + gd_size);
		head.overhead = metadata_size;
		head.unclean = 0;
		strcpy(head.endline_chars, VMDK_ENDLINECHARS);
	}
};

struct VMDKDesc
{
	VMDK_UINT version;
	VMDK_UINT cid;
	VMDK_UINT parentcid;
	VMDK_UINT hwver;
	string createtype;
	string adapter;
	string parentfile;
	
	VMDKDesc() 
	: version(1), cid(rand()), parentcid(VMDK_CIDNOPARENT), 
	  hwver(4), createtype("monolithicSparse"), adapter("lsilogic")
	{ }
	
	VMDKDesc(const Opts &opts) 
	: version(1), cid(rand()), parentcid(VMDK_CIDNOPARENT), 
	  hwver(opts.hwver), createtype("monolithicSparse"), adapter(opts.adapter) 
	{ }
	
	bool hasParent() const { return parentcid != VMDK_CIDNOPARENT && !parentfile.empty(); }

	void NewCID() { cid = rand(); }

	string GenDescriptor(const VMDKAttrs &attrs) const
	{
		char buf[80];
		string desc = "# Disk DescriptorFile\n"
					  "version=1\n";
		sprintf(buf, "%08x", cid);
		desc += "CID=" + string(buf) + "\n";
		sprintf(buf, "%08x", parentcid);
		desc += "parentCID=" + string(buf) + "\n";
		desc += "createType=\"" + createtype + "\"\n";
		
		if (!parentfile.empty())
			desc += "parentFileNameHint=\"" + parentfile + "\"\n";
		
		return desc + "\n";
	}

	string GenDDB(const VMDKAttrs &attrs, bool ddb) const
	{
		char buf[80];
		string desc = "# The Disk Data Base\n"
					  "#DDB\n\n";
		if (ddb)
		{
			sprintf(buf, "%u", hwver);
			desc += "ddb.virtualHWVersion = \"" + string(buf) + "\"\n";
			sprintf(buf, "%llu", attrs.getDataSize() / (255 * 63));
			desc += "ddb.geometry.cylinders = \"" + string(buf) + "\"\n";
			desc += "ddb.geometry.heads = \"255\"\n"
					"ddb.geometry.sectors = \"63\"\n"
					"ddb.adapterType = \"" + adapter + "\"\n";
		}
		return desc;
	}

	string GenExtDesc(const VMDKAttrs &attrs, const string &filename) const
	{
		char buf[80];
		string desc = "# Extent description\n";
		sprintf(buf, "%llu", attrs.getDataSize());
		desc += "RW " + string(buf) + " SPARSE \"" + filename + "\"\n\n";
		return desc + "\n";
	}
	
	string GenExtDesc(const VMDKAttrs &attrs, const VMDKFileRefList &refs) const
	{
		char buf[80];
		string desc = "# Extent description\n";
		VMDKFileRefList::const_iterator it = refs.begin();
		for (; it != refs.end(); it ++)
		{
			sprintf(buf, " %llu ", it->capacity);
			desc += it->access + string(buf) + it->type + " \"" + it->filename + "\"";
			if (it->offset != 0)
			{
				sprintf(buf, " %llu", it->offset);
				desc += buf;
			}
			desc += "\n";
		}
		return desc + "\n";
	}

	string GenDesc(const VMDKAttrs &attrs, const string &filename, bool ddb) const
	{
		string desc = GenDescriptor(attrs);
		desc += GenExtDesc(attrs, filename);
		desc += GenDDB(attrs, ddb);
		return desc;
	}
	
	string GenDesc(const VMDKAttrs &attrs, const VMDKFileRefList &refs, bool ddb) const
	{
		string desc = GenDescriptor(attrs);
		desc += GenExtDesc(attrs, refs);
		desc += GenDDB(attrs, ddb);
		return desc;
	}
};

static void fnsplit(const char *fullname, string &path, string &fn)
{
	const char *p = fullname + strlen(fullname);
	while (p > fullname && *p != '/' && *p != '\\') p --;
	if (*p == '/' || *p == '\\') p ++;
	fn = p; path.assign(fullname, p - fullname);
}

static bool isabspath(const char *filename)
{
	return filename[0] == '/' || filename[0] == '\\' ||
				(strlen(filename) > 3 && filename[1] == ':' &&
					(filename[2] == '/' || filename[2] == '\\'));
}

class FileReader
{
public:
	FileReader(int fd) : m_fd(fd), m_end(false)
	{
		m_pos = tell(m_fd);
	}
	
	char readbyte()
	{
		char ch = 0;
		if (!m_end)
		{
			if (read(m_fd, &ch, 1) != 1) ch = 0;
			if (ch == 0) m_end = true;
		}
		return ch;
	}
	
	void rewind()
	{
		lseek(m_fd, m_pos, SEEK_SET);
		m_end = false;
	}

private:
	int m_fd;
	long m_pos;
	bool m_end;
};

class DescFileParser
{
public:
	DescFileParser(FileReader *reader) 
	: m_reader(reader), m_status(stReady) {}

	size_t Parse()
	{
		char ch;
		size_t len;
		for (len = 0; (ch = m_reader->readbyte()) != 0; len ++)
		{
			if (ch == '\r') continue;
			switch (m_status)
			{
			case stReady:
				if (ch == '#') m_status = stComment;
				else if (ch == '"') m_status = stQuote;
				else if (ch != '\n' && !isspace(ch)) { m_token = ch; m_status = stToken; }
				break;
			case stToken:
				if (ch == '=') { shift(); m_token = "="; shift(); }
				else if (ch == '\n' || isspace(ch)) shift();
				else m_token += ch;
				break;
			case stQuote:
				if (ch == '"' || ch == '\n') shift();
				else m_token += ch;
				break;
			case stComment:
				if (ch == '\n') m_status = stReady;
				break;
			}
			
			if (ch == '\n' && !m_tokens.empty()) process();
		}
		if (!m_token.empty()) shift();
		if (!m_tokens.empty()) process();

		return len;
	}
	
	VMDKDesc getDesc() const { return m_desc; }
	VMDKFileRefList getRefFiles() const { return m_reffiles; }

private:
	FileReader *m_reader;
	VMDKDesc m_desc;
	enum { stReady, stToken, stQuote, stComment } m_status;
	string m_token;
	list<string> m_tokens;
	VMDKFileRefList m_reffiles;
	
	void shift()
	{
		if (!m_token.empty())
		{
			m_tokens.push_back(m_token);
			m_token.clear();
		}
		m_status = stReady;
	}
	
	void process()
	{
		string key = m_tokens.front(); m_tokens.pop_front();
		if (!m_tokens.empty() && m_tokens.front().compare("=") == 0)
		{
			string val;
			m_tokens.pop_front();
			if (!m_tokens.empty()) val = m_tokens.front();
			
			if (key.compare("version") == 0) m_desc.version = atoi(val.c_str());
			else if (key.compare("CID") == 0) sscanf(val.c_str(), "%x", &m_desc.cid);
			else if (key.compare("parentCID") == 0) sscanf(val.c_str(), "%x", &m_desc.parentcid);
			else if (key.compare("createType") == 0) m_desc.createtype = val;
			else if (key.compare("parentFileNameHint") == 0) m_desc.parentfile = val;
			else if (key.compare("ddb.virtualHWVersion") == 0) m_desc.hwver = atoi(val.c_str());
			else if (key.compare("ddb.adapterType") == 0) m_desc.adapter = val;
		}
		else if (m_tokens.size() >= 3)
		{
			VMDKFileRef fr;
			fr.access = key;
			sscanf(m_tokens.front().c_str(), "%llu", &fr.capacity); m_tokens.pop_front();
			fr.type = m_tokens.front(); m_tokens.pop_front();
			fr.filename = m_tokens.front(); m_tokens.pop_front();
			if (!m_tokens.empty()) sscanf(m_tokens.front().c_str(), "%llu", &fr.offset);
			m_reffiles.push_back(fr);
		}
		m_tokens.clear();
	}
};

class VMDKReader
{
public:
	VMDKReader(int fd) : m_fd(fd) { }
	virtual ~VMDKReader() { if (m_fd >= 0) close(m_fd); }
	
	int getFD() const { return m_fd; }

	VMDK_UINT getCID() const { return m_desc.cid; }
	const VMDKAttrs& getAttrs() const { return m_attrs; }
	const VMDKDesc& getDesc() const { return m_desc; }
	VMDK_SECT getCapacity() const { return m_capacity; }
	size_t getDiskSize(char unit) const
	{
		VMDK_SECT div = 2LL * 1024LL * 1024LL;
		if (unit == 'M') div /= 1024L;
		return m_capacity / div;
	}
	VMDK_SECT getGrainCount() const { return m_attrs.grain_count; }
	VMDK_SECT getDataGrainCount() const { return m_attrs.data_grain_count; }
	size_t getGrainBytes() const { return VMDK_GRAIN_SIZE * VMDK_SEC_SIZE; }
	
	virtual bool isSplitted() const = 0;
	virtual size_t ReadGrain(VMDK_SECT num, void *grain, size_t maxsz) = 0;
	
protected:
	int m_fd;
	VMDK_SECT m_capacity;
	VMDKDesc m_desc;
	VMDKAttrs m_attrs;
	
	size_t ParseDesc(long off, VMDKFileRefList *refs)
	{
		lseek(m_fd, off * VMDK_SEC_SIZE, SEEK_SET);
		FileReader reader(m_fd);
		DescFileParser parser(&reader);
		size_t len = parser.Parse();
		m_desc = parser.getDesc();
		if (refs) *refs = parser.getRefFiles();
		return len;
	}
};

typedef vector<VMDKReader*>	VMDKReaders;

class VMDKSingleFileReader : public VMDKReader
{
public:
	VMDKSingleFileReader(int fd)
	: VMDKReader(fd)
	{	
		lseek(m_fd, 0, SEEK_SET);
		read(m_fd, &m_head, sizeof(m_head));
		m_attrs.CalcAll(m_head.capacity, m_head.desc_off ? m_head.desc_size : 0);
		if (m_head.desc_off && m_head.desc_size)
			ParseDesc(m_head.desc_off, NULL);
		m_capacity = m_head.capacity;
	}

	bool isSplitted() const { return false; }

	size_t ReadGrain(VMDK_SECT num, void *grain, size_t maxsz)
	{
		if (num >= getGrainCount()) return 0;
		VMDK_SECT gt = num / m_head.gte_per_gt;
		VMDK_SECT id = num % m_head.gte_per_gt;
		VMDK_UINT gtoff, goff;
		lseek64(m_fd, (off_t)m_head.gd_off * VMDK_SEC_SIZE + gt * sizeof(gtoff), SEEK_SET);
		read(m_fd, &gtoff, sizeof(gtoff));
		if (gtoff == 0) return 0;
		lseek64(m_fd, (off_t)gtoff * VMDK_SEC_SIZE + id * sizeof(goff), SEEK_SET);
		read(m_fd, &goff, sizeof(goff));
		if (goff == 0) return 0;
		lseek64(m_fd, (off_t)goff * VMDK_SEC_SIZE, SEEK_SET);
		size_t sz = (size_t)(m_head.grain_size * VMDK_SEC_SIZE);
		if (sz > maxsz) sz = maxsz;
		int rd = read(m_fd, grain, sz);
		return rd > 0 ? (size_t)rd : 0;
	}

private:
	VMDK_SPEXTHEAD m_head;
};

typedef list<VMDKSingleFileReader*> SingleReaderList;

class VMDKMultiFileReader : public VMDKReader
{
public:
	VMDKMultiFileReader(int fd, const string &path)
	: VMDKReader(fd)
	{
		size_t len = ParseDesc(0, &m_refs);
	
		VMDKFileRefList::iterator it = m_refs.begin();
		for (m_capacity = 0; it != m_refs.end(); it ++)
		{
			string filename;
			if (isabspath(it->filename.c_str()))
				filename = it->filename;
			else
				filename = path + it->filename;
			fd = open(filename.c_str(), OPEN_RFLAGS);
			if (fd < 0) { cerr << "Failed to open " << filename << endl; exit(1); }
			VMDKSingleFileReader *sfr = new VMDKSingleFileReader(fd);
			m_files.push_back(sfr);
			m_capacity += sfr->getCapacity();
		}
		
		m_attrs.CalcAll(m_capacity, len);
	}

	~VMDKMultiFileReader()
	{
		SingleReaderList::iterator it = m_files.begin();
		for (; it != m_files.end(); it ++)
			delete *it;
	}
	
	bool isSplitted() const { return true; }
	
	size_t ReadGrain(VMDK_SECT num, void *grain, size_t maxsz)
	{
		VMDK_SECT base = 0;
		SingleReaderList::iterator it = m_files.begin();
		for (; it != m_files.end(); it ++)
		{
			if (num >= base && num < base + (*it)->getDataGrainCount())
				return (*it)->ReadGrain(num - base, grain, maxsz);
			base += (*it)->getDataGrainCount();
		}
		return 0;
	}

	const VMDKFileRefList& getFileRefs() const { return m_refs; }

private:
	VMDKFileRefList m_refs;
	SingleReaderList m_files;
};

class VMDKChainReader : public VMDKReader
{
public:
	VMDKChainReader(const VMDKReaders &readers, const string &path)
	: VMDKReader(-1), m_readers(readers), m_path(path), m_levels(Opts::CHAIN_ALL)
	{
		m_attrs = m_readers.back()->getAttrs();
		m_desc = m_readers.back()->getDesc();
		m_desc.parentcid = VMDK_CIDNOPARENT;
		m_desc.parentfile.clear();
		m_capacity = m_readers.back()->getCapacity();
	}
	
	~VMDKChainReader()
	{
		for (VMDKReaders::iterator it = m_readers.begin(); it != m_readers.end(); it ++)
			delete *it;
	}
	
	bool isSplitted() const { return m_readers[0]->isSplitted(); }
	
	size_t ReadGrain(VMDK_SECT num, void *grain, size_t maxsz)
	{
		VMDKReaders::iterator it = m_readers.begin();
		for (int n = 0; it != m_readers.end() && (m_levels == Opts::CHAIN_ALL || n < m_levels); it ++, n ++)
		{
			size_t sz = (*it)->ReadGrain(num, grain, maxsz);
			if (sz) return sz;
		}

		return 0;
	}

	int getLevels() const { return m_levels; }

	void setLevels(int levels) 
	{
		m_levels = levels; 
		if (m_levels != Opts::CHAIN_ALL && m_levels > 0 && m_levels < (int)m_readers.size())
		{
			m_desc.parentcid = m_readers[m_levels - 1]->getDesc().parentcid;
			m_desc.parentfile = m_readers[m_levels - 1]->getDesc().parentfile;
		}
		else
		{
			m_desc.parentcid = VMDK_CIDNOPARENT;
			m_desc.parentfile.clear();
		}
	}
	
	const VMDKReaders& getReaders() const { return m_readers; }

private:
	VMDKReaders m_readers;
	string m_path;
	int m_levels;
};

class VMDKOpener
{
public:
	static VMDKReader* Open(const string &filename, int flags = OPEN_RFLAGS)
	{
		int fd = open(filename.c_str(), flags);
		if (fd < 0) { cerr << "Failed to open " << filename << endl; exit(1); }
		
		VMDK_SPEXTHEAD head;
		read(fd, &head, 32);
		if (head.magic == VMDK_SPEXTMAGIC && head.version == VMDK_VERSION)
		{
			return new VMDKSingleFileReader(fd);
		}
		else
		{
			string path, file;
			fnsplit(filename.c_str(), path, file);
			return new VMDKMultiFileReader(fd, path);
		}
	}
	
	static VMDKChainReader* OpenChain(VMDKReader *reader, const char *fn)
	{
		if (!reader->getDesc().hasParent())
			return NULL;

		VMDKReaders readers;
		readers.push_back(reader);

		string path, file;
		fnsplit(fn, path, file);

		while (reader->getDesc().hasParent())
		{
			VMDKDesc desc = reader->getDesc();
			string filename = isabspath(desc.parentfile.c_str()) ? desc.parentfile : (path + desc.parentfile);
			reader = Open(filename);
			readers.push_back(reader);
		}

		return new VMDKChainReader(readers, path);
	}
};

class VMDKWriter
{
public:
	VMDKWriter(int fd, const VMDKAttrs &attrs) : m_fd(fd), m_attrs(attrs) {}
	virtual ~VMDKWriter() { Close(); }
	
	const VMDKAttrs& getAttrs() const { return m_attrs; }
	virtual void WriteGrain(VMDK_SECT num, const void *grain, size_t sz) = 0;
	virtual void Close() { if (m_fd >= 0) { close(m_fd); m_fd = -1; } }
	
protected:
	int m_fd;
	VMDKAttrs m_attrs;
};

class VMDKSingleFileWriter : public VMDKWriter
{
public:
	VMDKSingleFileWriter(int fd, const VMDKAttrs &attrs) 
	: VMDKWriter(fd, attrs) 
	{
		attrs.ExportHead(m_head);
		m_cursize = (VMDK_UINT)m_head.overhead;
	}
	
	void WriteGrain(VMDK_SECT num, const void *grain, size_t sz)
	{
		if (num < m_attrs.data_grain_count)
		{
			VMDK_SECT gt = num / m_head.gte_per_gt;
			VMDK_SECT id = num % m_head.gte_per_gt;
			lseek64(m_fd, (off_t)(m_head.rgd_off + m_attrs.gd_size + VMDK_GT_SIZE * gt) * VMDK_SEC_SIZE + 
							id * sizeof(VMDK_UINT), SEEK_SET);
			write(m_fd, &m_cursize, sizeof(m_cursize));
			lseek64(m_fd, (off_t)(m_head.gd_off + m_attrs.gd_size + VMDK_GT_SIZE * gt) * VMDK_SEC_SIZE + 
							id * sizeof(VMDK_UINT), SEEK_SET);
			write(m_fd, &m_cursize, sizeof(m_cursize));
			lseek64(m_fd, (off_t)m_cursize * VMDK_SEC_SIZE, SEEK_SET);
			size_t gsz = (size_t)(m_head.grain_size * VMDK_SEC_SIZE);
			if (sz > gsz) sz = gsz;
			write(m_fd, grain, sz);
			if (sz < gsz)
			{
				char *buf = new char[gsz - sz];
				memset(buf, 0, gsz - sz);
				write(m_fd, buf, gsz - sz);
				delete []buf;
			}
			m_cursize += m_head.grain_size;
		}
	}
	
private:
	VMDK_SPEXTHEAD m_head;
	VMDK_UINT m_cursize;
};

typedef list<VMDKSingleFileWriter*> SingleFileWriterList;

class VMDKMultiFileWriter : public VMDKWriter
{
public:
	VMDKMultiFileWriter(int fd, const VMDKAttrs &attrs, const SingleFileWriterList &sf)
	: VMDKWriter(fd, attrs), m_sfiles(sf)
	{}
	
	~VMDKMultiFileWriter()
	{
		Close();
		SingleFileWriterList::iterator it = m_sfiles.begin();
		for (; it != m_sfiles.end(); it ++)
			delete *it;
	}
	
	void Close()
	{
		SingleFileWriterList::iterator it = m_sfiles.begin();
		for (; it != m_sfiles.end(); it ++)
			(*it)->Close();
		VMDKWriter::Close();
	}

	void WriteGrain(VMDK_SECT num, const void *grain, size_t sz)
	{
		VMDK_SECT base = 0;
		SingleFileWriterList::iterator it = m_sfiles.begin();
		for (; it != m_sfiles.end(); it ++)
		{
			if (num >= base && num < base + (*it)->getAttrs().data_grain_count)
			{
				(*it)->WriteGrain(num - base, grain, sz);
				break;
			}
			base += (*it)->getAttrs().data_grain_count;
		}
	}

private:
	SingleFileWriterList m_sfiles;
};

class VMDKCreator
{
public:
	VMDKCreator(const char *fn, const VMDKAttrs &attrs, const VMDKDesc &desc)
	: m_filename(fn), m_attrs(attrs), m_desc(desc)
	{ 
		fnsplit(fn, m_path, m_file);
	}
	
	VMDKWriter* CreateSingleFile(bool ddb)
	{
		string descstr = m_desc.GenDesc(m_attrs, m_file, ddb);
		int fd = CreateSparseExtentFile(m_filename.c_str(), m_attrs, descstr.c_str(), descstr.length());
		return new VMDKSingleFileWriter(fd, m_attrs);
	}
	
	VMDKWriter* CreateSplittedFiles(bool ddb)
	{
		m_desc.createtype = "twoGbMaxExtentSparse";
		VMDK_SECT capacity = m_attrs.getDataSize(), size;
		VMDKFileRefList refs;
		SingleFileWriterList sfiles;
		string base, ext;
		size_t pos = m_file.rfind('.');
		if (pos == string::npos) base = m_file;
		else
		{
			base = m_file.substr(0, pos);
			ext = m_file.substr(pos);
		}
		
		char buf[80];
		for (pos = 1; capacity > 0; pos ++)
		{
			sprintf(buf, "-s%03u", pos);
			size = capacity;
			if (size > VMDK_SPLIT_SIZE) size = VMDK_SPLIT_SIZE;
			VMDKFileRef fr;
			fr.capacity = size;
			fr.filename = base + buf + ext;
			refs.push_back(fr);
			VMDKAttrs attrs(size, 0);
			int fd = CreateSparseExtentFile((m_path + fr.filename).c_str(), attrs, NULL, 0);
			VMDKSingleFileWriter *sfw = new VMDKSingleFileWriter(fd, attrs);
			sfiles.push_back(sfw);
			capacity -= size;
		}
		string descstr = m_desc.GenDesc(m_attrs, refs, ddb);
		int fd = open(m_filename.c_str(), OPEN_FLAGS, OPEN_MODE);
		if (fd < 0) { cerr << "Failed to create " << m_filename << endl; exit(1); }
		write(fd, descstr.c_str(), descstr.length());
		return new VMDKMultiFileWriter(fd, m_attrs, sfiles);
	}

private:
	string m_filename, m_path, m_file;
	VMDKAttrs m_attrs;
	VMDKDesc m_desc;

	static int CreateSparseExtentFile(const char *fn, const VMDKAttrs &attrs, const char *desc, size_t desc_size)
	{
		char sector[VMDK_SEC_SIZE];
		VMDK_SECT n, m, i, cnt = 1;
		
		int fd = open(fn, OPEN_FLAGS, OPEN_MODE);
		if (fd < 0) { cerr << "Failed to create file " << fn << endl; exit(1); }
	
		VMDK_SPEXTHEAD head;
		attrs.ExportHead(head);
		
		memset(sector, 0, sizeof(sector));
		memcpy(sector, &head, sizeof(head));
		write(fd, sector, sizeof(sector));
	
		for (n = 0; n < attrs.desc_size; n ++)
		{
			size_t len = desc ? desc_size : 0;
			if (len > sizeof(sector)) len = sizeof(sector);
			memset(sector, 0, sizeof(sector));
			if (len) memcpy(sector, desc, len);
			write(fd, sector, sizeof(sector));
			cnt ++;
			if (len)
			{
				desc += len;
				desc_size -= len;
			}
		}
		
		for (n = m = 0; n < attrs.gd_size; n ++)
		{
			memset(sector, 0, sizeof(sector));
			VMDK_UINT *p = (VMDK_UINT*)sector;
			for (i = 0; m < attrs.gt_num && i < (sizeof(sector) >> 2); i ++, m ++)
				p[i] = head.rgd_off + attrs.gd_size + m * VMDK_GT_SIZE;
			write(fd, sector, sizeof(sector));
			cnt ++;
		}
		memset(sector, 0, sizeof(sector));
		for (n = 0; n < attrs.gt_num * VMDK_GT_SIZE; n ++, cnt ++)
			write(fd, sector, sizeof(sector));
		
		for (n = m = 0; n < attrs.gd_size; n ++)
		{
			memset(sector, 0, sizeof(sector));
			VMDK_UINT *p = (VMDK_UINT*)sector;
			for (i = 0; m < attrs.gt_num && i < (sizeof(sector) >> 2); i ++, m ++)
				p[i] = head.gd_off + attrs.gd_size + m * VMDK_GT_SIZE;
			write(fd, sector, sizeof(sector));
			cnt ++;
		}
		memset(sector, 0, sizeof(sector));
		for (n = 0; n < attrs.gt_num * VMDK_GT_SIZE; n ++, cnt ++)
			write(fd, sector, sizeof(sector));
		
		cnt &= VMDK_GRAIN_SIZE - 1;
		if (cnt)
			for (; cnt < VMDK_GRAIN_SIZE; cnt ++)
				write(fd, sector, sizeof(sector));
	
		return fd;
	}
};

class CommandProcessor
{
public:
	CommandProcessor(Opts *opts) : m_opts(opts) { }
	virtual ~CommandProcessor() { }
	virtual int ProcessCommand(int argc, char *argv[]) = 0;
	
protected:
	Opts *m_opts;
};

class CommandInfo : public CommandProcessor
{
public:
	CommandInfo(Opts *opts) : CommandProcessor(opts) { }
	int ProcessCommand(int argc, char *argv[])
	{
		if (argc < 2) usage();
		VMDKReader *reader = VMDKOpener::Open(argv[1]);
		VMDKChainReader *cr = NULL;
		VMDKAttrs attrs = reader->getAttrs();
		VMDKDesc desc = reader->getDesc();
		if (desc.hasParent())
		{
			cr = VMDKOpener::OpenChain(reader, argv[1]);
			attrs = cr->getAttrs();
			desc = cr->getDesc();
		}

		cout << "Version        : " << desc.version << endl
			 << "CreateType     : " << desc.createtype << endl
			 << "HardwareVersion: " << desc.hwver << endl
			 << "AdapterType    : " << desc.adapter << endl
			 << "Capacity       : " << reader->getCapacity() << " Sectors, " << reader->getDiskSize('G') << " GBytes" << endl
		     << "Grains         : " << reader->getDataGrainCount() << " (" << reader->getGrainBytes() << " bytes each)" << endl
			 << "GrainTables    : " << attrs.gt_num << endl
			 << "GrainDirSize   : " << attrs.gd_size << endl
			 << "Overhead       : " << attrs.metadata_size << endl
		     << endl;
		
		if (cr)
		{
			cout << "Chained VMDKs:" << endl;
			cout << "000 " << argv[1] << endl;
			const VMDKReaders &rds = cr->getReaders();
			char buf[80];
			for (size_t cnt = 0; cnt < rds.size() - 1; cnt ++)
			{
				sprintf(buf, "%03u ", cnt + 1);
				cout << buf << rds[cnt]->getDesc().parentfile << endl;
			}
			cout << endl;
		}
	
		delete reader;
		return 0;
	}
};

class CommandConvert : public CommandProcessor
{
public:
	CommandConvert(Opts *opts) : CommandProcessor(opts) { }
	int ProcessCommand(int argc, char *argv[])
	{
		if (argc < 3) usage();
		VMDKReader *reader = VMDKOpener::Open(argv[1], OPEN_WFLAGS);
		const VMDKAttrs attrs = reader->getAttrs();
		VMDKDesc desc = reader->getDesc();
		if (attrs.desc_size == 0)
		{
			cerr << "No description found in file, please convert the root file in cloned chain." << endl;
			return 1;
		}
		else if (desc.hasParent())
		{
			cerr << "This is not a root file in cloned chain, please convert the root file in cloned chain." << endl;
			return 1;
		}
		else
		{
			for (int i = 2; i < argc; i ++)
			{
				char *p = strchr(argv[i], '=');
				if (p)
				{
					if (strncmp(argv[i], "hwver=", 6) == 0) desc.hwver = atoi(p + 1);
					else if (strncmp(argv[i], "adapter=", 8) == 0) desc.adapter = p + 1;
				}
			}
			
			if (reader->isSplitted())
			{
				string str = desc.GenDesc(attrs, ((VMDKMultiFileReader*)reader)->getFileRefs(), true);
				delete reader;
				int fd = open(argv[1], OPEN_FLAGS, OPEN_MODE);
				if (fd < 0) { cerr << "Failed to re-write file " << argv[1] << endl; exit(1); }
				write(fd, str.c_str(), str.length());
				close(fd);
				return 0;
			}
			else
			{
				string path, file;
				fnsplit(argv[1], path, file);
				char *buf = new char[attrs.desc_size * VMDK_SEC_SIZE];
				memset(buf, 0, attrs.desc_size * VMDK_SEC_SIZE);
				strcpy(buf, desc.GenDesc(attrs, file, true).c_str());
				lseek64(reader->getFD(), VMDK_SEC_SIZE, SEEK_SET);
				write(reader->getFD(), buf, attrs.desc_size * VMDK_SEC_SIZE);
				delete []buf;
			}
		}
		delete reader;
		return 0;
	}
};

class CommandCreate : public CommandProcessor
{
public:
	CommandCreate(Opts *opts) : CommandProcessor(opts) { }
	int ProcessCommand(int argc, char *argv[])
	{
		if (argc < 2 || m_opts->disk_size == 0) usage();
		if (m_opts->disk_size < 8 * 1024L * 1024L)
		{
			cerr << "Disk size must be greater than 8M" << endl;
			exit(1);
		}
	
		VMDKAttrs attrs(*m_opts, 512);
		VMDKDesc desc(*m_opts);
		
		VMDKCreator creator(argv[1], attrs, desc);
		VMDKWriter *writer;
		
		if (m_opts->split && attrs.getDataSize() > VMDK_SPLIT_SIZE)
			writer = creator.CreateSplittedFiles(true);
		else
			writer = creator.CreateSingleFile(true);
		
		writer->Close();
		delete writer;
	
		return 0;
	}
};

class CommandClone : public CommandProcessor
{
public:
	CommandClone(Opts *opts) : CommandProcessor(opts) { }
	int ProcessCommand(int argc, char *argv[])
	{
		if (argc < 3) usage();
		VMDKReader *reader = VMDKOpener::Open(argv[1]);
		VMDKAttrs attrs = reader->getAttrs();
		VMDKDesc desc = reader->getDesc();
		desc.NewCID();
		desc.parentcid = reader->getCID();
		desc.parentfile = argv[1];
		VMDKCreator creator(argv[2], attrs, desc);
		VMDKWriter *writer;
		if ((reader->isSplitted() && !m_opts->merge) || m_opts->split) 
			writer = creator.CreateSplittedFiles(false);
		else 
			writer = creator.CreateSingleFile(false);
		writer->Close();
		delete writer;
		delete reader;
		return 0;
	}
};

class CommandCopy : public CommandProcessor
{
public:
	CommandCopy(Opts *opts) : CommandProcessor(opts) { }
	int ProcessCommand(int argc, char *argv[])
	{
		if (argc < 3) usage();
		VMDKReader *reader = VMDKOpener::Open(argv[1]);
		if (m_opts->chains != 0)
		{
			VMDKChainReader *cr = VMDKOpener::OpenChain(reader, argv[1]);
			if (cr)
			{
				reader = cr;
				cr->setLevels(m_opts->chains);
			}
		}

		VMDKAttrs attrs = reader->getAttrs();
		VMDKDesc desc = reader->getDesc();
		desc.NewCID();
		VMDKCreator creator(argv[2], attrs, desc);
		VMDKWriter *writer;
		if ((reader->isSplitted() && !m_opts->merge) || m_opts->split) 
			writer = creator.CreateSplittedFiles(!desc.hasParent());
		else 
			writer = creator.CreateSingleFile(!desc.hasParent());
		
		VMDK_SECT grains, g;
		size_t sz, gsz = reader->getGrainBytes();
		char *grain = new char[gsz];
		grains = reader->getGrainCount();
		cout << "Copying " << grains << " Grains ..." << endl;
		for (g = 0; g < grains; g ++)
		{
			cout << (g + 1) << "/" << grains << ", " << ((g + 1) * 100 / grains) << "%\r";
			sz = reader->ReadGrain(g, grain, gsz);
			if (sz > 0)
				writer->WriteGrain(g, grain, sz);
		}
		cout << endl;
		delete []grain;
		writer->Close();
		delete writer;
		delete reader;
		return 0;
	}
};

static string parse_cmd(int &argc, char *argv[])
{
	string cmd;
	if (argc >= 2)
	{
		if (argv[1][0] == '-')
		{
			char *p = argv[0] + strlen(argv[0]);
			while (p > argv[0] && *p != '/' && *p != '\\') p --;
			if (*p == '/' || *p == '\\') p ++;
			cmd = p;
		}
		else
		{
			cmd = argv[1];
			int i;
			for (i = 2; i < argc; i ++)
				argv[i - 1] = argv[i];
			argv[i] = NULL;
			argc --;
		}
	}
	return cmd;
}

static unsigned long long parse_size(char *str)
{
	unsigned long long factor = 0, val;
	int len = strlen(str);
	if (str[len - 1] == 'G') factor = 1024L;
	else if (str[len - 1] == 'M') factor = 1L;
	else usage();
	str[len - 1] = 0;
	val = atoi(str);
	if (val == 0) usage();
	return val * factor * 1024L * 1024L;
}

static void parse_opts(int &argc, char *argv[], Opts &opts)
{
	int i, j;
	for (i = 1; i < argc; i ++)
	{
		if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--size") == 0)
		{
			if (++ i < argc) opts.disk_size = parse_size(argv[i]);
			else usage();
		}
		else if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--adapter") == 0)
		{
			if (++ i < argc) opts.adapter = argv[i];
			else usage();
		}
		else if (strcmp(argv[i], "-r") == 0 || strcmp(argv[i], "--hwver") == 0)
		{
			if (++ i < argc)
			{
				opts.hwver = atoi(argv[i]);
				if (opts.hwver == 0) usage();
			}
			else usage();
		}
		else if (strcmp(argv[i], "-c") == 0 || strcmp(argv[i], "--chains") == 0)
		{
			if (++ i < argc)
			{
				if (strcmp(argv[i], "all") == 0) opts.chains = Opts::CHAIN_ALL;
				else
				{
					opts.chains = atoi(argv[i]);
					if (opts.chains <= 0) usage();
				}
			}
			else usage();
		}
		else if (strcmp(argv[i], "-p") == 0 || strcmp(argv[i], "--split") == 0)
			opts.split = true;
		else if (strcmp(argv[i], "-m") == 0 || strcmp(argv[i], "--merge") == 0)
			opts.merge = true;
		else if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0)
			usage();
		else if (argv[i][0] == '-') usage();
		else break;
	}
	
	argc -= i;
	for (j = 1; j <= argc; j ++, i ++)
		argv[j] = argv[i];
	argv[j] = NULL;
	argc = j;
}

int main(int argc, char *argv[])
{
	string cmd = parse_cmd(argc, argv);
	Opts opts;
	parse_opts(argc, argv, opts);
	
	srand(time(NULL));
	
	map<string, CommandProcessor*> cmds;
	cmds["info"]   = new CommandInfo(&opts);
	cmds["convert"] = new CommandConvert(&opts);
	cmds["create"] = new CommandCreate(&opts);
	cmds["clone"]  = new CommandClone(&opts);
	cmds["copy"]   = new CommandCopy(&opts);
	
	map<string, CommandProcessor*>::iterator it = cmds.find(cmd);
	if (it == cmds.end()) usage();
	return it->second->ProcessCommand(argc, argv);
}
