/*
Copyright (C) 2006 Evan Teran
                   eteran@alum.rit.edu

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., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

#include "DialogFunctions.h"
#include "DebuggerCoreInterface.h"
#include "Debugger.h"
#include "MemoryRegions.h"
#include "Util.h"
#include "Configuration.h"
#include "Instruction.h"
#include "State.h"
#include "ArchProcessorInterface.h"

#include <QHeaderView>
#include <QMessageBox>
#include <QVector>
#include <QSortFilterProxyModel>
#include <QList>

#define MIN_REFCOUNT 2

//------------------------------------------------------------------------------
// Name: DialogFunctions(QWidget *parent)
// Desc: 
//------------------------------------------------------------------------------
DialogFunctions::DialogFunctions(QWidget *parent) : QDialog(parent) {
	ui.setupUi(this);
	ui.tableView->verticalHeader()->hide();
	ui.tableView->horizontalHeader()->setStretchLastSection(true);
	ui.tableView->horizontalHeader()->setResizeMode(QHeaderView::ResizeToContents);
	
	ui.tableWidget->verticalHeader()->hide();
	ui.tableWidget->horizontalHeader()->setStretchLastSection(true);
	ui.tableWidget->horizontalHeader()->setResizeMode(QHeaderView::ResizeToContents);
	
	m_FilterModel = new QSortFilterProxyModel(this);
	connect(ui.txtSearch, SIGNAL(textChanged(const QString &)), m_FilterModel, SLOT(setFilterFixedString(const QString &)));
}

//------------------------------------------------------------------------------
// Name: on_tableWidget_cellDoubleClicked(int row, int column)
// Desc: follows the found item in the data view
//------------------------------------------------------------------------------
void DialogFunctions::on_tableWidget_cellDoubleClicked(int row, int column) {
	Q_UNUSED(column);
		
	QTableWidgetItem *const item = ui.tableWidget->item(row, 0);
	const edb::address_t addr = item->data(Qt::UserRole).toULongLong();
	edb::v1::jumpToAddress(addr);
}

//------------------------------------------------------------------------------
// Name: showEvent(QShowEvent *)
// Desc: 
//------------------------------------------------------------------------------
void DialogFunctions::showEvent(QShowEvent *) {
	m_FilterModel->setFilterKeyColumn(3);
	m_FilterModel->setSourceModel(&edb::v1::memoryRegions());
	ui.tableView->setModel(m_FilterModel);
	
	ui.progressBar->setValue(0);
	ui.tableWidget->setRowCount(0);
}

//------------------------------------------------------------------------------
// Name: isThunk(edb::address_t address) const
// Desc: basically returns true if the first instruction of the function is a 
//       jmp
//------------------------------------------------------------------------------
bool DialogFunctions::isThunk(edb::address_t address) const {
	quint8 buf[Instruction::MAX_SIZE];
	int buf_size = sizeof(buf);
	bool ret = false;
	if(edb::v1::getInstructionBytes(address, buf, buf_size)) {
		const Instruction insn(buf, buf_size, address, std::nothrow);
		ret = insn.valid() && insn.type() == Instruction::OP_JMP;
	}
	
	return ret;
}

//------------------------------------------------------------------------------
// Name: findFunctionCalls()
// Desc: 
//------------------------------------------------------------------------------
void DialogFunctions::findFunctionCalls() {

	const QItemSelectionModel *const selModel = ui.tableView->selectionModel();
	const QModelIndexList sel = selModel->selectedRows();
	
	if(sel.size() == 0) {
		QMessageBox::information(this, "No Region Selected", "You must select a region which is to be scanned for functions.");
		return;
	}
	
	const edb::address_t pageSize = edb::v1::debuggerBase->pageSize();

	foreach(const QModelIndex selected_item, sel) {

		const QModelIndex index = m_FilterModel->mapToSource(selected_item);
		const MemoryRegions::Region *const region = reinterpret_cast<const MemoryRegions::Region *>(index.internalPointer());
		
		// do the search for this region!	
		if(region != 0) {

			const edb::address_t size_in_pages = region->size() / pageSize;
			quint8 *pages = new quint8[size_in_pages * pageSize];
			
			if(edb::v1::debuggerBase->readPages(region->start, pages, size_in_pages)) {
				for(edb::address_t i = 0; i < size_in_pages * pageSize; ++i) {
					const Instruction insn(&pages[i], region->size() - i, region->start + i, std::nothrow);
					
					if(insn.valid() && insn.type() == Instruction::OP_CALL) {

						State state;
						state.setInstructionPointer(region->start + i);

						const Operand &op = insn.operand(0);
						
						if(op.generalType() == Operand::TYPE_REL) {
							edb::address_t ea = edb::v1::archProcessor()->getEffectiveAddress(op, state);
							
							// skip over ones which are : call <label>; label:
							if(ea != state.instructionPointer() + insn.size()) {
								if(ea >= region->start && ea < region->end) {
									
									m_ResultSet[ea].entry_address	= ea;
									m_ResultSet[ea].end_address		= ea;
									m_ResultSet[ea].reference_count++;	
								}
							}
						}

						ui.progressBar->setValue(CalculatePercentage(i, region->size()));
					}
				}
			}
			
			releaseArray(pages);
			
		}
	}
}

//------------------------------------------------------------------------------
// Name: findFunctionEnd(edb::address_t start_address, edb::address_t end_address, int &update_count)
// Desc: 
//------------------------------------------------------------------------------
edb::address_t DialogFunctions::findFunctionEnd(edb::address_t start_address, edb::address_t end_address, int &update_count) {
	edb::address_t ret = start_address;
	QList<edb::address_t>		jump_targets;
	QMap<edb::address_t, int>	code_addresses;
	
	update_count = 0;
	
	// we start with the entry point of the function
	jump_targets.push_front(start_address);
	
	// while no more jump targets... (includes entry point)
	while(!jump_targets.empty()) {
		edb::address_t addr = jump_targets.front();
		jump_targets.pop_front();

		// keep going until we go out of bounds	
		while(addr >= start_address && addr < end_address) {
				
			quint8 buf[Instruction::MAX_SIZE];
			int buf_size = sizeof(buf);
			if(!edb::v1::getInstructionBytes(addr, buf, buf_size)) {
				break;
			}
			
			const Instruction insn(buf, buf_size, addr, std::nothrow);
			if(!insn.valid()) {
				break;
			}

			// ok, it was a valid instruction, let's add it to the
			// list of possible "code addresses"
			code_addresses.insert(addr, insn.size());

			State state;
			state.setInstructionPointer(addr);

			// is it a branch?
			if(insn.type() == Instruction::OP_JCC) {

				const Operand &op = insn.operand(0);
				if(op.generalType() == Operand::TYPE_REL) {
					const edb::address_t ea = edb::v1::archProcessor()->getEffectiveAddress(op, state);

					if(!code_addresses.contains(ea)) {
						if(!jump_targets.contains(ea)) {
							jump_targets.push_front(ea);
						}
					}
				}
			} else if(insn.type() == Instruction::OP_JMP) {

				const Operand &op = insn.operand(0);
				if(op.generalType() == Operand::TYPE_REL) {
					const edb::address_t ea = edb::v1::archProcessor()->getEffectiveAddress(op, state);
					
					// is it a jump to another function's entry point?
					// if so, this is a dead end, resume from other branches
					// but give the target a bonus reference
					funcmap_t::iterator it = m_ResultSet.find(ea);
					if(it != m_ResultSet.end()) {
						if(it.value().reference_count < MIN_REFCOUNT) {
							it.value().reference_count++;
							++update_count;
							break;
						}
					}
					
					
					if(!code_addresses.contains(ea)) {
						addr = ea;
						continue;
					} else {
						break;
					}
				}
			} else if(insn.type() == Instruction::OP_CALL) {

				const Operand &op = insn.operand(0);
				if(op.generalType() == Operand::TYPE_REL) {
					const edb::address_t ea = edb::v1::archProcessor()->getEffectiveAddress(op, state);

					// skip over ones which are : call <label>; label:
					if(ea != state.instructionPointer() + insn.size()) {
					
						funcmap_t::iterator it = m_ResultSet.find(ea);
						if(it != m_ResultSet.end()) {
							if(it.value().reference_count < MIN_REFCOUNT) {
								it.value().reference_count++;
								++update_count;
							}
						}					
					}
				} else if(op.generalType() == Operand::TYPE_EXPRESSION) {
				#if 0
					printf("found jump table? %08x\n", addr);
				#endif
				}
			} else if(insn.type() == Instruction::OP_RET) {
				break;

			}
			
			addr += insn.size();
		}
	}
	
	// find the last entry
	for(QMap<edb::address_t, int>::const_iterator it = code_addresses.begin(); it != code_addresses.end(); ++it) {
		const edb::address_t a = it.key() + it.value() - 1;
		ret = qMax(a, ret);
	}
	
	return ret;
}

//------------------------------------------------------------------------------
// Name: findFunctionEnds()
// Desc: 
//------------------------------------------------------------------------------
int DialogFunctions::findFunctionEnds() {

	int ret = 0;
	funcmap_t::iterator it = m_ResultSet.begin();
	while(it != m_ResultSet.end()) {
		FunctionInfo &info = it.value();
		
		const funcmap_t::iterator next = ++it;
		
		if(info.reference_count >= MIN_REFCOUNT) {
			if(next != m_ResultSet.end()) {
				int updates;
				const edb::address_t end = findFunctionEnd(info.entry_address, next.value().entry_address, updates);
				ret += updates;
				info.end_address = end;
			} else {
				int updates;
				const edb::address_t end = findFunctionEnd(info.entry_address, info.entry_address | 0xfff, updates);
				ret += updates;
				info.end_address = end;
			}
		}	
		
		it = next;
	}
	return ret;
}

//------------------------------------------------------------------------------
// Name: findStackFrames()
// Desc: 
//------------------------------------------------------------------------------
void DialogFunctions::findStackFrames() {

	// give bonus if we see a:
	// push ebp; mov ebp, esp;

	for(funcmap_t::iterator it = m_ResultSet.begin(); it != m_ResultSet.end(); ++it) {
	
		bool bonus = false;
	
		FunctionInfo &info = it.value();
		edb::address_t addr = info.entry_address;
		
		quint8 buf[Instruction::MAX_SIZE];
		int buf_size = sizeof(buf);
		if(!edb::v1::getInstructionBytes(addr, buf, buf_size)) {
			continue;
		}
			
		const Instruction insn(buf, buf_size, addr, std::nothrow);
		if(!insn.valid()) {
			continue;
		}
		
		if(insn.type() == Instruction::OP_PUSH) {
			const Operand &op = insn.operand(0);
			if(op.completeType() == Operand::TYPE_REGISTER) {
				if(op.reg == Operand::REG_EBP) {
					bonus = true;
				}
			}
		}
		
		if(bonus) {
			bonus = false;
			addr += insn.size();
			
			int buf_size = sizeof(buf);
			if(!edb::v1::getInstructionBytes(addr, buf, buf_size)) {
				continue;
			}

			const Instruction insn(buf, buf_size, addr, std::nothrow);
			if(!insn.valid()) {
				continue;
			}

			if(insn.type() == Instruction::OP_MOV) {
				const Operand &op0 = insn.operand(0);
				const Operand &op1 = insn.operand(1);
				if(op0.completeType() == Operand::TYPE_REGISTER && op1.completeType() == Operand::TYPE_REGISTER) {
					if(op0.reg == Operand::REG_EBP && op1.reg == Operand::REG_ESP) {
						bonus = true;
					}
				}
			}
		}
		
		if(bonus) {
			info.reference_count++;
		}
	}
}

//------------------------------------------------------------------------------
// Name: bonusMain()
// Desc: 
//------------------------------------------------------------------------------
void DialogFunctions::bonusMain() {
	const QStringList args = edb::v1::getProcessArguments();
	if(!args.isEmpty()) {
		const QString s = args[0];

		edb::address_t main = edb::v1::locateMainFunction(s);


		// give bonus if we see a:
		// push ebp; mov ebp, esp;

		for(funcmap_t::iterator it = m_ResultSet.begin(); it != m_ResultSet.end(); ++it) {

			FunctionInfo &info = it.value();
			FunctionInfo &info_next = (it + 1).value();


			if(main >= info.entry_address && main < info_next.entry_address) {
				m_ResultSet[main].entry_address		= main;
				m_ResultSet[main].end_address		= main;
				m_ResultSet[main].reference_count	= 2;
			}
		}
	}
}

//------------------------------------------------------------------------------
// Name: doFind()
// Desc: 
//------------------------------------------------------------------------------
void DialogFunctions::doFind() {

	QString temp;
	
	m_ResultSet.clear();

	edb::v1::log("[Function Finder] finding obvious function calls...");
	findFunctionCalls();
	
	edb::v1::log("[Function Finder] attempting to add 'main' to the list...");
	bonusMain();
	
	edb::v1::log("[Function Finder] finding stack frames...");
	findStackFrames();
	
	int updates;
	do {
		edb::v1::log("[Function Finder] calculating function bounds...");
		updates = findFunctionEnds();
		edb::v1::log("[Function Finder] got %d updates", updates);
	} while(updates != 0);
	
	ui.tableWidget->setRowCount(0);		
	ui.tableWidget->setSortingEnabled(false);

	for(funcmap_t::const_iterator it = m_ResultSet.begin(); it != m_ResultSet.end(); ++it) {
	
		const FunctionInfo &info = it.value();

		const int row = ui.tableWidget->rowCount();
		ui.tableWidget->insertRow(row);
		
		// entry point
		QTableWidgetItem *const p = new QTableWidgetItem(temp.sprintf(EDB_FMT_PTR, info.entry_address));
		p->setData(Qt::UserRole, info.entry_address);
		ui.tableWidget->setItem(row, 0, p);

		// upper bound of the function
		if(info.reference_count >= MIN_REFCOUNT) {
			ui.tableWidget->setItem(row, 1, new QTableWidgetItem(temp.sprintf(EDB_FMT_PTR, info.end_address)));
		}

		// reference count
		QTableWidgetItem *const itemCount = new QTableWidgetItem;
		itemCount->setData(Qt::DisplayRole, info.reference_count);
		ui.tableWidget->setItem(row, 2, itemCount);
		
		// type
		if(isThunk(it.key())) {
			ui.tableWidget->setItem(row, 3, new QTableWidgetItem("Thunk"));
		} else {
			ui.tableWidget->setItem(row, 3, new QTableWidgetItem("Standard Function"));
		}
	}

	ui.tableWidget->setSortingEnabled(true);

}

//------------------------------------------------------------------------------
// Name: on_btnFind_clicked()
// Desc: 
//------------------------------------------------------------------------------
void DialogFunctions::on_btnFind_clicked() {
	ui.btnFind->setEnabled(false);
	ui.progressBar->setValue(0);
	doFind();
	ui.progressBar->setValue(100);
	ui.btnFind->setEnabled(true);
}
