/*
 * Copyright (c) 2001 Jason Tishler
 *
 * 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 2
 * 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, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 *
 * $Id: rebase.cc,v 1.5 2001/12/18 20:21:19 jtishler Exp $
 */

#include <iostream>
#include <sstream>
#include <string>
#include <time.h>
#include <stdlib.h>
#include <getopt.h>
#include <sys/cygwin.h>
#include <windows.h>
#include <imagehlp.h>

using namespace std;

string PosixToWin32(const string& aPosixPath);
void ParseArgs(int argc, char* argv[]);
unsigned long StringToUlong(const string& aString);
void Usage();

ULONG theImageBase = 0;
BOOL theDownFlag = FALSE;
ULONG theOffset = 0;
int theArgsIndex = 0;

void
main(int argc, char* argv[])
{
	ParseArgs(argc, argv);
	ULONG aNewImageBase = theImageBase;

	for (int i = theArgsIndex; i < argc; i++)
	{
		if (theDownFlag)
			aNewImageBase -= theOffset;

		string aFile = PosixToWin32(argv[i]);
		ULONG anOldImageSize, anOldImageBase, aNewImageSize;
		ULONG aPrevNewImageBase = aNewImageBase;
		BOOL aStatus = ReBaseImage(
			const_cast<char*>(aFile.c_str()), // CurrentImageName
			0, // SymbolPath
			TRUE, // fReBase
			FALSE, // fRebaseSysfileOk
			theDownFlag, // fGoingDown
			0, // CheckImageSize
			&anOldImageSize, // OldImageSize
			&anOldImageBase, // OldImageBase
			&aNewImageSize, // NewImageSize
			&aNewImageBase, // NewImageBase
			time(0)); // TimeStamp

		// ReBaseImage seems to never returns false!
		DWORD aStatus2 = GetLastError();
		if (aStatus2 != 0)
		{
			cerr << "ReBaseImage() failed with last error = " <<
				GetLastError() << endl;
			exit(2);
		}
		cout << aFile << hex << ": new base = " <<
			((theDownFlag) ? aNewImageBase : aPrevNewImageBase) <<
			", new size = " << aNewImageSize + theOffset << endl;

		if (!theDownFlag)
			aNewImageBase += theOffset;
	}

	exit(0);
}

string
PosixToWin32(const string& aPosixPath)
{
	char aWin32Path[MAX_PATH];
	cygwin_conv_to_win32_path(aPosixPath.c_str(), aWin32Path);
	return aWin32Path;
}

void
ParseArgs(int argc, char* argv[])
{
	const char* anOptions = "b:do:";
	for (int anOption; (anOption = getopt(argc, argv, anOptions)) != -1;)
	{
		switch (anOption)
		{
			case 'b':
				theImageBase = StringToUlong(optarg);
				break;
			case 'd':
				theDownFlag = TRUE;
				break;
			case 'o':
				theOffset = StringToUlong(optarg);
				break;
			default:
				Usage();
				exit(1);
				break;
		}
	}

	if (theImageBase == 0)
	{
		Usage();
		exit(1);
	}

	theArgsIndex = optind;
}

unsigned long
StringToUlong(const string& aString)
{
	stringstream ss;
	unsigned long aUlong;
	string::size_type anIndex = aString.find("0x");
	if (anIndex == 0)
		ss << hex << string(aString, 2, aString.size() - 2);
	else
		ss << aString;
	ss >> aUlong;
	return aUlong;
}

void
Usage()
{
	cerr << "usage: rebase -b BaseAddress [-d] -o Offset ImageFileName ..." << endl;
}