/****************************************************************************
 *                                                                          *
 *   Copyright (C) 2005  Michael Buesch <mbuesch@freenet.de>                *
 *                                                                          *
 *   This program is free software; you can redistribute it and/or modify   *
 *   it under the terms of the GNU General Public License version 2         *
 *   as published by the Free Software Foundation.                          *
 *                                                                          *
 *   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. Additionally a copy of the                    *
 *   GNU General Public License is available here:                          *
 *   http://passwordmanager.sourceforge.net/gplv2.txt                       *
 *   http://www.gnu.org/licenses/gpl.txt                                    *
 *                                                                          *
 ****************************************************************************/

#include <qapplication.h>
#include <qlabel.h>
#include <qlineedit.h>
#include <qtextedit.h>
#include <qpushbutton.h>
#include <qclipboard.h>
#include <qmessagebox.h>
#include <qvaluelist.h>
#include <qfiledialog.h>
#include <qfile.h>
#include <qdatastream.h>

#include "regtrack.h"
#include "stack.h"
#include "advanced.h"



/***************************************************************************
 *  class RegTrackContext
 ***************************************************************************/

RegTrackContext::RegTrackContext(RegTrack *_rt)
 : rt (_rt)
 , _current (0)
{
	/* Add an initial context.
	 * We must _always_ have at least one context.
	 */
	QByteArray *buf = new QByteArray;
	QDataStream ds(*buf, IO_WriteOnly);
	rt->saveContext(ds);
	blobs.append(buf);
	names.append("Default");
}

RegTrackContext::~RegTrackContext()
{
	clear();
}

int RegTrackContext::count() const
{
	return blobs.count();
}

int RegTrackContext::current() const
{
	return _current;
}

QString RegTrackContext::name(int index) const
{
	return names[index];
}

int RegTrackContext::cloneContext(int index, const QString &name)
{
	QByteArray *buf = new QByteArray;

	saveContext(index);
	buf->duplicate(*(blobs[index]));
	blobs.append(buf);
	names.append(name);
	return blobs.count() - 1;
}

bool RegTrackContext::deleteContext(int index)
{
	int newIndex;

	if (blobs.count() == 1)
		return false;
	if (index == current()) {
		for (newIndex = 0; newIndex < count(); newIndex++) {
			if (newIndex != index)
				break;
		}
		switchContext(newIndex, false);
	}
	delete blobs[index];
	blobs.remove(blobs.at(index));
	names.remove(names.at(index));
	if (_current > index)
		_current--;
	return true;
}

bool RegTrackContext::renameContext(int index,
				    const QString &newName)
{
	if (index < 0 || static_cast<unsigned int>(index) >= names.count())
		return false;
	names[index] = newName;
	return true;
}

void RegTrackContext::loadContext(int index)
{
	QByteArray *buf = blobs[index];
	QDataStream ds(*buf, IO_ReadOnly);
	rt->loadContext(ds);
}

void RegTrackContext::saveContext(int index)
{
	QByteArray *buf = blobs[index];
	QDataStream ds(*buf, IO_WriteOnly);
	rt->saveContext(ds);
}

void RegTrackContext::switchContext(int index, bool saveCurrent)
{
	if (index == _current)
		return;
	if (saveCurrent && _current >= 0)
		saveContext(_current);
	_current = index;
	loadContext(index);
}

void RegTrackContext::clear()
{
	QValueList<QByteArray *>::iterator i = blobs.begin(),
					   end = blobs.end();
	for ( ; i != end; ++i)
		delete *i;
	blobs.clear();
	names.clear();
	_current = -1;
}

int RegTrackContext::appendEmptyContext(const QString &name)
{
	QByteArray *buf = new QByteArray;
	blobs.append(buf);
	names.append(name);
	return blobs.count() - 1;
}


/***************************************************************************
 *  class RegTrack
 ***************************************************************************/

QFont RegTrack::dataFont ("Courier");

RegTrack::RegTrack(QWidget *parent,
		   const char *name,
		   WFlags f)
 : QWidget (parent, name, f)
 , advancedWidget (0)
{
	QHBoxLayout *h;
	QHBoxLayout *h2;
	QVBoxLayout *v;
	QLabel *label;

	layout = new QVBoxLayout(this);

	// GPR first line
	h = new QHBoxLayout(layout);
	label = new QLabel("EAX:", this);
	label->setFont(dataFont);
	h->addWidget(label);
	eax = new QLineEdit(this);
	h->addWidget(eax);
	eaxClear = new QPushButton("R", this);
	h->addWidget(eaxClear);
	label = new QLabel("EBX:", this);
	label->setFont(dataFont);
	h->addWidget(label);
	ebx = new QLineEdit(this);
	h->addWidget(ebx);
	ebxClear = new QPushButton("R", this);
	h->addWidget(ebxClear);

	// GPR second line
	h = new QHBoxLayout(layout);
	label = new QLabel("ECX:", this);
	label->setFont(dataFont);
	h->addWidget(label);
	ecx = new QLineEdit(this);
	h->addWidget(ecx);
	ecxClear = new QPushButton("R", this);
	h->addWidget(ecxClear);
	label = new QLabel("EDX:", this);
	label->setFont(dataFont);
	h->addWidget(label);
	edx = new QLineEdit(this);
	h->addWidget(edx);
	edxClear = new QPushButton("R", this);
	h->addWidget(edxClear);

	// GPR third line
	h = new QHBoxLayout(layout);
	label = new QLabel("ESI:", this);
	label->setFont(dataFont);
	h->addWidget(label);
	esi = new QLineEdit(this);
	h->addWidget(esi);
	esiClear = new QPushButton("R", this);
	h->addWidget(esiClear);
	label = new QLabel("EDI:", this);
	label->setFont(dataFont);
	h->addWidget(label);
	edi = new QLineEdit(this);
	h->addWidget(edi);
	ediClear = new QPushButton("R", this);
	h->addWidget(ediClear);

	// Stack browser
	freeSize = NULL;
	h = new QHBoxLayout(layout);
	stack = new StackView(this);
	h->addWidget(stack);

	// EBP/ESP
	v = new QVBoxLayout(h);
	h2 = new QHBoxLayout(v);
	label = new QLabel("EBP:", this);
	label->setFont(dataFont);
	h2->addWidget(label);
	ebp = new QLineEdit(this);
	h2->addWidget(ebp);
	ebpClear = new QPushButton("R", this);
	h2->addWidget(ebpClear);
	h2 = new QHBoxLayout(v);
	label = new QLabel("ESP:", this);
	label->setFont(dataFont);
	h2->addWidget(label);
	esp = new QLineEdit(this);
	esp->setReadOnly(true);
	h2->addWidget(esp);

	// Stack allocation
	v->addStretch();
	h2 = new QHBoxLayout(v);
	allocSize = new HexSpinBox(4, 2048, 4, this);
	h2->addWidget(allocSize);
	stackAlloc = new QPushButton("<-- Stack Alloc", this);
	h2->addWidget(stackAlloc);
	stackPush = new QPushButton("PUSH", this);
	h2->addWidget(stackPush);
	h2 = new QHBoxLayout(v);
	freeSize = new HexSpinBox(4, 2048, 4, this);
	freeSize->setMaxValue(sizeof(uint32_t));
	h2->addWidget(freeSize);
	stackFree = new QPushButton("<-- Stack Free", this);
	h2->addWidget(stackFree);
	stackPop = new QPushButton("POP", this);
	h2->addWidget(stackPop);

	// Other buttons
	v->addStretch();
	h2 = new QHBoxLayout(v);
	reset = new QPushButton("Reset", this);
	h2->addWidget(reset);
	advanced = new QPushButton("Advanced", this);
	h2->addWidget(advanced);
	advanced->setToggleButton(true);

	// Notes
	bottomLayout = new QHBoxLayout(layout);
	notes = new QTextEdit(this);
	bottomLayout->addWidget(notes);
	advancedWidget = new AdvancedWidget(this);
	bottomLayout->addWidget(advancedWidget);
	advancedWidget->hide();

	// Fonts / size
	eax->setFont(dataFont);
	ebx->setFont(dataFont);
	ecx->setFont(dataFont);
	edx->setFont(dataFont);
	esi->setFont(dataFont);
	edi->setFont(dataFont);
	ebp->setFont(dataFont);
	esp->setFont(dataFont);
	notes->setFont(dataFont);
	resize(QSize(650, 350));
	setCaption("regtrack v" VERSION " - i386 Register/Stack tracker");

	// Connections
	connect(stackAlloc, SIGNAL(clicked()),
		this, SLOT(stackAllocClicked()));
	connect(stackFree, SIGNAL(clicked()),
		this, SLOT(stackFreeClicked()));
	connect(stackPush, SIGNAL(clicked()),
		this, SLOT(stackPushClicked()));
	connect(stackPop, SIGNAL(clicked()),
		this, SLOT(stackPopClicked()));
	connect(reset, SIGNAL(clicked()),
		this, SLOT(resetClicked()));
	connect(advanced, SIGNAL(toggled(bool)),
		this, SLOT(advancedClicked(bool)));
	connect(eaxClear, SIGNAL(clicked()),
		eax, SLOT(clear()));
	connect(ebxClear, SIGNAL(clicked()),
		ebx, SLOT(clear()));
	connect(ecxClear, SIGNAL(clicked()),
		ecx, SLOT(clear()));
	connect(edxClear, SIGNAL(clicked()),
		edx, SLOT(clear()));
	connect(ebpClear, SIGNAL(clicked()),
		ebp, SLOT(clear()));
	connect(esiClear, SIGNAL(clicked()),
		esi, SLOT(clear()));
	connect(ediClear, SIGNAL(clicked()),
		edi, SLOT(clear()));

	context = new RegTrackContext(this);
	update();
}

RegTrack::~RegTrack()
{
	delete context;
}

void RegTrack::update()
{
	int32_t curEsp;
	QString tmp;

	curEsp = -(stack->size() - sizeof(int32_t));
	tmp = QString::number(curEsp, 16).right(8);
	while (tmp.length() < 8)
		tmp = QString("0") + tmp;
	esp->setText(QString("0x") + tmp);
}

void RegTrack::stackAllocClicked()
{
	int size;
	int nr;

	size = allocSize->value();
	nr = size / sizeof(int32_t);
	for ( ; nr; nr--)
		stack->push(QString());
	update();
}

void RegTrack::stackFreeClicked()
{
	int size;
	int nr;

	size = freeSize->value();
	nr = size / sizeof(int32_t);
	for ( ; nr; nr--)
		stack->pop();
	update();
}

void RegTrack::stackPushClicked()
{
	QString str;

	str = clipboardGet().left(100);
	str.replace('\n', ' ');
	str.replace('\r', ' ');
	stack->push(str);
	update();
}

void RegTrack::stackPopClicked()
{
	clipboardPut(stack->pop());
	update();
}

void RegTrack::resetClicked()
{
	int res;

	res = QMessageBox::question(this, "Reset Stack?",
				    "Reset Stack?",
				    QMessageBox::Yes,
				    QMessageBox::No);
	if (res == QMessageBox::Yes)
		stack->reset();
	update();
	res = QMessageBox::question(this, "Reset Registers?",
				    "Reset Registers?",
				    QMessageBox::Yes,
				    QMessageBox::No);
	if (res == QMessageBox::Yes) {
		eax->clear();
		ebx->clear();
		ecx->clear();
		esi->clear();
		edi->clear();
		edx->clear();
		ebp->clear();
	}
	update();
}

void RegTrack::advancedClicked(bool on)
{
	if (on)
		advancedWidget->show();
	else
		advancedWidget->hide();
}

QString RegTrack::clipboardGet() const
{
	QClipboard *cb;
	QString str;

	cb = QApplication::clipboard();
	if (cb->supportsSelection())
		str = cb->text(QClipboard::Selection);
	else
		str = cb->text(QClipboard::Clipboard);

	return str;
}

void RegTrack::clipboardPut(const QString &str) const
{
	QClipboard *cb;

	cb = QApplication::clipboard();
	if (cb->supportsSelection())
		cb->setText(str, QClipboard::Selection);
	cb->setText(str, QClipboard::Clipboard);
}

void RegTrack::loadContext(QDataStream &ds)
{
	QString str;

	/* This is where _all_ context data is loaded from a
	 * data stream.
	 */
	ds >> str;
	eax->setText(str);
	ds >> str;
	ebx->setText(str);
	ds >> str;
	ecx->setText(str);
	ds >> str;
	edx->setText(str);
	ds >> str;
	esi->setText(str);
	ds >> str;
	edi->setText(str);
	stack->load(ds);
	ds >> str;
	ebp->setText(str);
	ds >> str;
	notes->setText(str);

	update();
}

void RegTrack::saveContext(QDataStream &ds)
{
	/* This is where _all_ context data is saved to a
	 * data stream.
	 */
	ds << eax->text();
	ds << ebx->text();
	ds << ecx->text();
	ds << edx->text();
	ds << esi->text();
	ds << edi->text();
	stack->save(ds);
	ds << ebp->text();
	ds << notes->text();
}

#define REGTRACK_MAGIC		(static_cast<uint32_t>(0x20051224))
#define REGTRACK_FILEVER	(static_cast<uint32_t>(2))

void RegTrack::loadFile()
{
	uint32_t tmp;
	QString str;
	QSize size;
	QPoint pos;
	int i, contextCount;
	int index;

	str = QFileDialog::getOpenFileName();
	if (str.isEmpty())
		return;

	QFile file(str);
	file.open(IO_ReadOnly);
	QDataStream ds(&file);

	ds >> tmp;
	if (tmp != REGTRACK_MAGIC) {
		QMessageBox::warning(this,
				     "regtrack: Invalid Fileformat",
				     "Invalid Fileformat.");
		return;
	}
	ds >> tmp;
	if (tmp != REGTRACK_FILEVER) {
		int res;
		res = QMessageBox::question(this,
					    "regtrack: Fileversion Mismatch",
					    QString(
					    "Fileversion Mismatch "
					    "(file: %1, expected: %2)\n"
					    "Do you want to continue?")
					    .arg(QString::number(tmp))
					    .arg(QString::number(REGTRACK_FILEVER)),
					    QMessageBox::Yes,
					    QMessageBox::No);
		if (res != QMessageBox::Yes)
			return;
	}
	ds >> size;
	resize(size);
	ds >> pos;
	move(pos);

	ds >> contextCount;
	if (contextCount > 0)
		context->clear();
	for (i = 0; i < contextCount; i++) {
		ds >> str;
		index = context->appendEmptyContext(str);
		loadContext(ds);
		context->saveContext(index);
	}

	file.close();
}

void RegTrack::saveFile()
{
	QString fn;
	int i, contextCount;

	fn = QFileDialog::getSaveFileName();
	if (fn.isEmpty())
		return;
	if (!fn.endsWith(".regtrack"))
		fn += ".regtrack";

	QFile file(fn);
	file.open(IO_ReadWrite);
	QDataStream ds(&file);

	ds << REGTRACK_MAGIC;
	ds << REGTRACK_FILEVER;
	ds << size();
	ds << pos();

	contextCount = context->count();
	ds << contextCount;
	for (i = 0; i < contextCount; i++) {
		context->switchContext(i);
		ds << context->name(i);
		saveContext(ds);
	}

	file.close();
}