/*
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 "SymbolManager.h"

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

#define MIN_REFCOUNT 2

#if defined(EDB_X86)
static const edb::Operand::Register STACK_REG = edb::Operand::REG_ESP;
static const edb::Operand::Register FRAME_REG = edb::Operand::REG_EBP;
#elif defined(EDB_X86_64)
static const edb::Operand::Register STACK_REG = edb::Operand::REG_RSP;
static const edb::Operand::Register FRAME_REG = edb::Operand::REG_RBP;
#endif


//------------------------------------------------------------------------------
// 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[edb::Instruction::MAX_SIZE];
	int buf_size = sizeof(buf);
	bool ret = false;
	if(edb::v1::getInstructionBytes(address, buf, buf_size)) {
		const edb::Instruction insn(buf, buf_size, address, std::nothrow);
		ret = insn.valid() && insn.type() == edb::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 edb::Instruction insn(&pages[i], region->size() - i, region->start + i, std::nothrow);
					
					if(insn.valid() && insn.type() == edb::Instruction::OP_CALL) {

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

						const edb::Operand &op = insn.operand(0);
						
						if(op.generalType() == edb::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()));
					}
				}
			}
			
			release_array(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[edb::Instruction::MAX_SIZE];
			int buf_size = sizeof(buf);
			if(!edb::v1::getInstructionBytes(addr, buf, buf_size)) {
				break;
			}
			
			const edb::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);
			
			const edb::Instruction::Type type = insn.type();

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

				const edb::Operand &op = insn.operand(0);
				if(op.generalType() == edb::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(type == edb::Instruction::OP_JMP) {

				const edb::Operand &op = insn.operand(0);
				if(op.generalType() == edb::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(type == edb::Instruction::OP_CALL) {

				const edb::Operand &op = insn.operand(0);
				if(op.generalType() == edb::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() == edb::Operand::TYPE_EXPRESSION) {
				#if 0
					printf("found jump table? %08x\n", addr);
				#endif
				}
			} else if(type == edb::Instruction::OP_RET) {
				break;
			} else if(type == edb::Instruction::OP_HLT) {
				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) {
			int updates;
			edb::address_t next_entry;
		
			if(next != m_ResultSet.end()) {
				next_entry = next.value().entry_address;
			} else {
				// just round up to next page, since we don't have a good guess for end...
				next_entry = ((info.entry_address + (0x1000 - 1)) & ~(0x1000 - 1));
				
			}
			const edb::address_t end = findFunctionEnd(info.entry_address, next_entry, 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[edb::Instruction::MAX_SIZE];
		
		unsigned int i = 0;
		
		while(i < 2) {
			// gets the bytes for the instruction
			int buf_size = sizeof(buf);
			if(!edb::v1::getInstructionBytes(addr, buf, buf_size)) {
				break;
			}
			
			// decode it
			const edb::Instruction insn(buf, buf_size, addr, std::nothrow);
			if(!insn.valid()) {
				break;;
			}
			
			// which part of the sequence are we looking at?
			switch(i++) {
			case 0:
				// are we looking at a "push ebp" ?
				if(insn.type() == edb::Instruction::OP_PUSH) {
					const edb::Operand &op = insn.operand(0);
					if(op.completeType() == edb::Operand::TYPE_REGISTER) {
						if(op.reg == FRAME_REG) {
							break;
						}
					}
				}
				++i;
				break;
			case 1:
				// are we looking at a "mov ebp, esp" ?
				if(insn.type() == edb::Instruction::OP_MOV) {
					const edb::Operand &op0 = insn.operand(0);
					const edb::Operand &op1 = insn.operand(1);
					if(op0.completeType() == edb::Operand::TYPE_REGISTER && op1.completeType() == edb::Operand::TYPE_REGISTER) {
						if(op0.reg == FRAME_REG && op1.reg == STACK_REG) {
							bonus = true;
						}
					}
				}
				break;
			default:
				break;
			}
						
			addr += insn.size();
		}
		
		if(bonus) {
			info.reference_count++;
		}
	}
}

//------------------------------------------------------------------------------
// Name: bonusMain()
// Desc: 
//------------------------------------------------------------------------------
void DialogFunctions::bonusMain() {
	const QString s = edb::v1::getProcessExecutable();
	if(!s.isEmpty()) {
		const edb::address_t main = edb::v1::locateMainFunction(s);
		
		if(main != 0x00000000) {


			// give bonus to main
			for(funcmap_t::iterator it = m_ResultSet.begin(); it != m_ResultSet.end(); ++it) {

				FunctionInfo &info				= it.value();
				const 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;
					break;
				}
			}
		}
	}
}

//------------------------------------------------------------------------------
// Name: bonusSymbols()
// Desc: 
//------------------------------------------------------------------------------
void DialogFunctions::bonusSymbols() {
	// give bonus if we have a symbol for the address

	for(funcmap_t::iterator it = m_ResultSet.begin(); it != m_ResultSet.end(); ++it) {
		const edb::address_t addr = it.value().entry_address;

		const SymbolManager::Symbol *const sym = edb::v1::symbolManager().findSymbol(addr);
		if(sym != 0) {

			m_ResultSet[addr].entry_address		= addr;
			m_ResultSet[addr].end_address		= addr;
			m_ResultSet[addr].reference_count++;
		}
	}
}

//------------------------------------------------------------------------------
// 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] attempting to add functions with symbols to the list...");
	bonusSymbols();
	
	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);
}
