mainwindow.cpp

#include "mainwindow.h"
#include "ui_mainwindow.h"
#include "pe-struct.h"
#include <assert.h>
#include <QtCore/QDebug>
#include <QtCore/QFile>
#include <QtCore/QSettings>
#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);
	ui->ReferenceBrowser->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
	if(config->value("SavingConfirm", true).toBool()) {
		if(QMessageBox::information(this, "PE Header Editor", tr("Are you sure to overwrite this file?"), QMessageBox::Yes | QMessageBox::No) != QMessageBox::Yes) {
			return false;
		}
	}
	QFile file(current_file, this);
	// If we use QFile::WriteOnly, the file will be truncated, that is a bug of Qt?
	if(!file.open(QFile::ReadWrite)) {
		QString msg = tr("Cannot open file %1 for writing").arg(current_file);
		QMessageBox::warning(this, 0, msg);
		ui->statusBar->showMessage(msg, 10);
		return false;
	}
	file.seek(pe_offset);
	QDataStream s(&file);
	if(s.writeRawData((const char *)pe, sizeof(PortableExecutable)) < 1) {
		QString msg = tr("Write failed");
		QMessageBox::warning(this, 0, msg);
		ui->statusBar->showMessage(msg, 10);
		return false;
	}
	file.close();
	ui->statusBar->showMessage(tr("File %1 saved").arg(current_file));
	return true;
}

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::init_value_reference() {
	QTreeWidget
}*/

void MainWindow::set_pe_header_item(const QString &name, uint8_t value, const QString &comment) {
	QString v;
	if(config->value("Hexadecimal", true).toBool()) v.sprintf("0x%hhx", value);
	else v.sprintf("%hhu", value);
	set_pe_header_item(name, v, comment, "uint8");
}

void MainWindow::set_pe_header_item(const QString &name, uint16_t value, const QString &comment) {
	QString v;
	if(config->value("Hexadecimal", true).toBool()) v.sprintf("0x%hx", value);
	else v.sprintf("%hu", value);
	set_pe_header_item(name, v, comment, "uint16");
}

void MainWindow::set_pe_header_item(const QString &name, uint32_t value, const QString &comment) {
	QString v;
	if(config->value("Hexadecimal", true).toBool()) v.sprintf("0x%x", value);
	else v.sprintf("%u", value);
	set_pe_header_item(name, v, comment, "uint32");
}

void MainWindow::set_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;
	QList<QTreeWidgetItem *> r = ui->TreeWidget->findItems(name, Qt::MatchFixedString, 0);
	int exists = r.length();
	if(exists) {
		assert(exists == 1);
		item = r[0];
	} else {
		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);
	if(!exists) ui->TreeWidget->addTopLevelItem(item);
}

QString MainWindow::get_comment_from_value_reference(const QTreeWidget *value_reference, const QString &value) {
	int value_count = value_reference->topLevelItemCount();
	if(!value_count) return QString();
	qDebug() << value;
	while(value_count--) {
		QTreeWidgetItem *value_item = value_reference->topLevelItem(value_count);
		if(value_item->text(0) == value) return value_item->text(1);
	}
	return QString();
}

QString MainWindow::get_comment_and_update(const QString &name, uint32_t value, bool update_struct, bool update_reference) {
	if(name == "machine") {
		if(update_struct) pe->machine = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Machine\nWhat platform the program running"));
			ui->stackedWidget->setCurrentWidget(ui->page_machine);
			return QString();
		}
		QString v;
		v.sprintf("0x%x", value);
		return get_comment_from_value_reference(ui->ValueReference_machine, v);
	} else if(name == "sections_count") {
		if(update_struct) pe->sections_count = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Sections Count"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString();
	} else if(name == "time_stamp") {
		if(update_struct) pe->time_stamp = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Time Stamp"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "symbol_table_pointer") {
		if(update_struct) pe->symbol_table_pointer = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Symbol Table Pointer"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "symbols_count") {
		if(update_struct) pe->symbols_count = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Symbols Count"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "optional_header_size") {
		if(update_struct) pe->optional_header_size = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Optional Header Size"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "characteristics") {
		if(update_struct) pe->characteristics = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Characteristics"));
			ui->stackedWidget->setCurrentWidget(ui->page_characteristics);
			return QString ();
		}
		QString v;
		v.sprintf("0x%x", value);
		return get_comment_from_value_reference(ui->ValueReference_characteristics, v);
	} else if(name == "magic") {
		if(update_struct) pe->magic = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("magic"));
			ui->stackedWidget->setCurrentWidget(ui->page_magic);
		}
		QString v;
		v.sprintf("0x%x", value);
		return get_comment_from_value_reference(ui->ValueReference_magic, v);
	} else if(name == "major_linker_version") {
		if(update_struct) pe->major_linker_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Major Linker Version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "minor_linker_version") {
		if(update_struct) pe->minor_linker_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Minor Linker Version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "code_size") {
		if(update_struct) pe->code_size = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Code Size"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "initialized_data_size") {
		if(update_struct) pe->initialized_data_size = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Initialized data size"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "uninitialized_data_size") {
		if(update_struct) pe->uninitialized_data_size = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Uninitialized data size"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "entry_point_address") {
		if(update_struct) pe->entry_point_address = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Entry point address"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "code_base") {
		if(update_struct) pe->code_base = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Code base"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "data_base") {
		if(update_struct) pe->data_base = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Data base"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "image_base") {
		if(update_struct) pe->image_base = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Image base"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "section_alignment") {
		if(update_struct) pe->section_alignment = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Section alignment"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "file_alignment") {
		if(update_struct) pe->file_alignment = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("File alignment"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "major_os_version") {
		if(update_struct) pe->major_os_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Major OS version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "minor_os_version") {
		if(update_struct) pe->minor_os_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Minor OS version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "major_image_version") {
		if(update_struct) pe->major_image_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Major image version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "minor_image_version") {
		if(update_struct) pe->minor_image_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Minor image version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "major_subsystem_version") {
		if(update_struct) pe->major_subsystem_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Major subsystem version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "minor_subsystem_version") {
		if(update_struct) pe->minor_subsystem_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Minor subsystem version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "may_be_a_version") {
		if(update_struct) pe->may_be_a_version = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Reserved, may be a version"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "image_size") {
		if(update_struct) pe->image_size = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Image size"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "headers_size") {
		if(update_struct) pe->headers_size = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Headers size"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "check_sum") {
		if(update_struct) pe->check_sum = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Check sum"));
			ui->stackedWidget->setCurrentWidget(ui->page_empty);
		}
		return QString ();
	} else if(name == "subsystem") {
		if(update_struct) pe->subsystem = value;
		if(update_reference) {
			ui->ReferenceBrowser->setText(tr("Subsystem"));
			ui->stackedWidget->setCurrentWidget(ui->page_subsystem);
			return QString();
		}
		QString v;
		v.sprintf("%u", value);
		return get_comment_from_value_reference(ui->ValueReference_subsystem, v);
	}
	return QString();
}

/*
void MainWindow::update_references(const QString &name) {
	int items_count = ui->stackedWidget->count();
	for(int i=1; i<items_count; i++) {
	}
}*/

void MainWindow::update_item_value(QTreeWidgetItem *item) {
	//update_references(item->text(0));
	QString v = item->text(2);
	int base = v.startsWith("0x") ? 16 : 10;
	item->setText(3, get_comment_and_update(item->text(0), v.toUInt(NULL, base), true, false));
}

void MainWindow::update_references(QTreeWidgetItem *item) {
	// Didn't need value when we updating the reference
	get_comment_and_update(item->text(0), 0, false, true);
}

#define DEFINE_PE_DEADER_ITEM(NAME) set_pe_header_item(#NAME, pe->NAME, get_comment_and_update(#NAME, pe->NAME, false, false))
void MainWindow::update_table() {
	//set_pe_header_item("machine", pe->machine, get_comment("machine", pe->machine));
	//set_pe_header_item("subsystem", pe->subsystem, get_comment("subsystem", pe->subsystem));
	DEFINE_PE_DEADER_ITEM(machine);
	DEFINE_PE_DEADER_ITEM(sections_count);
	DEFINE_PE_DEADER_ITEM(time_stamp);
	DEFINE_PE_DEADER_ITEM(symbol_table_pointer);
	DEFINE_PE_DEADER_ITEM(symbols_count);
	DEFINE_PE_DEADER_ITEM(optional_header_size);
	DEFINE_PE_DEADER_ITEM(characteristics);
	DEFINE_PE_DEADER_ITEM(magic);
	DEFINE_PE_DEADER_ITEM(major_linker_version);
	DEFINE_PE_DEADER_ITEM(minor_linker_version);
	DEFINE_PE_DEADER_ITEM(code_size);
	DEFINE_PE_DEADER_ITEM(initialized_data_size);
	DEFINE_PE_DEADER_ITEM(uninitialized_data_size);
	DEFINE_PE_DEADER_ITEM(entry_point_address);
	DEFINE_PE_DEADER_ITEM(code_base);
	DEFINE_PE_DEADER_ITEM(data_base);
	DEFINE_PE_DEADER_ITEM(image_base);
	DEFINE_PE_DEADER_ITEM(section_alignment);
	DEFINE_PE_DEADER_ITEM(file_alignment);
	DEFINE_PE_DEADER_ITEM(major_os_version);
	DEFINE_PE_DEADER_ITEM(minor_os_version);
	DEFINE_PE_DEADER_ITEM(major_image_version);
	DEFINE_PE_DEADER_ITEM(minor_image_version);
	DEFINE_PE_DEADER_ITEM(major_subsystem_version);
	DEFINE_PE_DEADER_ITEM(minor_subsystem_version);
	DEFINE_PE_DEADER_ITEM(may_be_a_version);
	DEFINE_PE_DEADER_ITEM(image_size);
	DEFINE_PE_DEADER_ITEM(headers_size);
	DEFINE_PE_DEADER_ITEM(check_sum);
	DEFINE_PE_DEADER_ITEM(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);
}