mainwindow.cpp

#include "mainwindow.h"
#include "ui_mainwindow.h"
#include "pe-struct.h"
#include <QtCore/QDebug>
#include <QtCore/QFile>
#include <QtGui/QFileDialog>
#include <QtGui/QMessageBox>

MainWindow::MainWindow(QWidget *parent) :
QMainWindow(parent), ui(new Ui::MainWindow) {
	ui->setupUi(this);
	is_modified = false;
	pe = NULL;
	set_current_file(0);
}

MainWindow::~MainWindow() {
	delete ui;
}

void MainWindow::load(const QString &filename) {
	QFile file(filename, this);
	if(!file.open(QFile::ReadOnly)) {
		QString msg = tr("Load file %1 failed").arg(filename);
		QMessageBox::warning(this, 0, msg);
		ui->statusBar->showMessage(msg, 10);
		return;
	}
	file.seek(0x3c);
	QDataStream s(&file);
	if(s.readRawData((char *)&pe_offset, sizeof pe_offset) < 1) {
		QString msg = tr("Cannot locate the PE header");
		QMessageBox::warning(this, 0, msg);
		goto failed;
	}
	//qDebug("pe_offset = %hhu", pe_offset);
	if(!pe_offset) {
		QString msg = tr("Unexcepted PE offset, %1 may not a valid PE file").arg(filename);
		QMessageBox::warning(this, 0, msg);
		goto failed;
	}
	file.seek(pe_offset);
	if(!pe) pe = new PortableExecutable;
	if(s.readRawData((char *)pe, sizeof(PortableExecutable)) < 1) {
		QString msg = tr("Cannot read PE header");
		QMessageBox::warning(this, 0, msg);
		goto failed;
	}
	//qDebug("pe->pe_signature = 0x%x", pe->pe_signature);
	if(pe->pe_signature != 0x4550) {
		QString msg = tr("The PE signature not found, %1 may not a valid PE file\n").arg(filename);
		QMessageBox::warning(this, 0, msg);
		goto failed;
	}
	update_table();
	ui->TreeWidget->setEnabled(true);
	ui->action_SaveAs->setEnabled(true);
	set_current_file(filename);
	return;

failed:
	if(pe) {
		delete pe;
		pe = NULL;
	}
	file.close();
}

bool MainWindow::save_file(const QString &filename) {
	// TODO: copy current file to new file, and make changes to this new file by calling save_file
}

void MainWindow::open_file() {
	if(!maybe_save()) return;
	QString filename = QFileDialog::getOpenFileName(this, NULL, NULL, NULL, NULL, QFileDialog::DontUseNativeDialog);
	if(!filename.isEmpty()) load(filename);
}

bool MainWindow::save_file() {
	//if(is_untitled) return save_file_as();		// We don't want create a new file this time...
	// TODO: reading the config, may pop a message box
	//return save_file(current_file);			// Don't use this function, because that is too large when we saving changes to current pe file
	// TODO: saving changes to current file
}

bool MainWindow::save_file_as() {
	QString filename = QFileDialog::getSaveFileName(this, tr("Save As"), current_file, NULL, NULL, QFileDialog::DontUseNativeDialog);
	if(filename.isEmpty()) return false;
	return save_file(filename);
}

void MainWindow::modified() {
	setWindowModified(true);
	is_modified = true;
	ui->action_Save->setEnabled(true);
}

bool MainWindow::maybe_save() {
	if(is_modified) {
		QMessageBox::StandardButton r = QMessageBox::warning(this, 0,
			tr("The table has been modified.\nDo you want to save your changes?"),
			QMessageBox::Save | QMessageBox::Discard | QMessageBox::Cancel);
		if(r == QMessageBox::Save) return save_file();
		if(r == QMessageBox::Cancel) return false;
	}
	return true;
}

QString MainWindow::stripped_name(const QString &filename) {
	return QFileInfo(filename).fileName();
}

void MainWindow::set_current_file(const QString &filename) {
	//static int seq_number = 1;
	//is_untitled = filename.isEmpty();
	//if(is_untitled) {
	if(filename.isEmpty()) {
		setWindowTitle(tr("PE Header Editor"));
	} else {
		current_file = QFileInfo(filename).canonicalFilePath();
		setWindowTitle(QString("%1[*] - %2").arg(stripped_name(current_file)).arg(tr("PE Header Editor")));
	}

	is_modified = false;
	setWindowModified(false);
}

void MainWindow::add_pe_header_item(const QString &name, uint8_t value, const QString &comment) {
	QString v;
	v.sprintf("%hhu", value);
	add_pe_header_item(name, v, comment, "uint8");
}

void MainWindow::add_pe_header_item(const QString &name, uint16_t value, const QString &comment) {
	QString v;
	v.sprintf("%hu", value);
	add_pe_header_item(name, v, comment, "uint16");
}

void MainWindow::add_pe_header_item(const QString &name, uint32_t value, const QString &comment) {
	QString v;
	v.sprintf("%u", value);
	add_pe_header_item(name, v, comment, "uint32");
}

void MainWindow::add_pe_header_item(const QString &name, const QString &value, const QString &comment, const QString &type) {
	//qDebug("function: add_pe_header_item(%s, %s, %s, %s)", name.toLocal8Bit().data(), value.toLocal8Bit().data(), comment.toLocal8Bit().data(), type.toLocal8Bit().data());
	qDebug() << value;
	QTreeWidgetItem *item = new QTreeWidgetItem;
	item->setFlags(Qt::ItemIsSelectable | Qt::ItemIsEnabled);
	item->setText(0, name);
	item->setText(1, type);
	item->setText(2, value);
	item->setText(3, comment);
	ui->TreeWidget->addTopLevelItem(item);
}

QString MainWindow::get_comment(const QString &name, uint32_t value) {
	if(name == "machine") {
		switch(value) {
			case 0x14c:
				return "i386";
			case 0x1c0:
				return "ARM";
		}
	} else if(name == "subsystem") {
		switch(value) {
			case 1:
				return "Native";
			case 2:
				return "Windows (GUI)";
			case 3:
				return "Console";
			case 4:
				return "POSIX";
			case 9:
				return "Windows CE";
		}
	}
	return QString();
}

void MainWindow::update_comment(QTreeWidgetItem *item) {
	item->setText(3, get_comment(item->text(0), item->text(2).toUInt()));
}

void MainWindow::update_table() {
	add_pe_header_item("machine", pe->machine, get_comment("machine", pe->machine));
	add_pe_header_item("subsystem", pe->subsystem, get_comment("subsystem", pe->subsystem));
}

void MainWindow::edit_item(QTreeWidgetItem *item, int column) {
	qDebug("slot: edit_item(%p, %d)", item, column);
	item->setFlags(Qt::ItemIsSelectable | Qt::ItemIsEnabled | Qt::ItemIsEditable);
	ui->TreeWidget->editItem(item, 2);
	item->setFlags(Qt::ItemIsSelectable | Qt::ItemIsEnabled);
}