/*
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 "MemoryRegions.h"
#include "SymbolManager.h"
#include "DebuggerCoreInterface.h"
#include "Debugger.h"
#include "Util.h"
#include "State.h"
#include <sstream>
#include <iostream>
#include <string>
#include <asm/unistd.h>
#include <sys/mman.h>
#include <QApplication>
#include <QFile>
#include <QTextStream>
#include <QMessageBox>

const unsigned int MemoryRegions::Region::READ	= PROT_READ;
const unsigned int MemoryRegions::Region::WRITE	= PROT_WRITE;
const unsigned int MemoryRegions::Region::EXEC	= PROT_EXEC;

class BackupInfo {
public:
	BackupInfo(std::size_t size, edb::address_t address, uint8_t perms);
	~BackupInfo();
private:
	BackupInfo &operator=(const BackupInfo &);
	BackupInfo(const BackupInfo&);
	
public:
	uint32_t perms() const { return m_Perms; }
public:
	bool backup();
	bool restore();
private:
	uint8_t 					*m_Buffer;
	std::size_t					m_Size;
	edb::address_t		m_Address;
	uint32_t					m_Perms;
	State						m_State;
public:
	DebugEventHandlerInterface	*eventhandler;
};

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
BackupInfo::BackupInfo(std::size_t size, edb::address_t address, uint8_t perms) : m_Size(size), m_Address(address), m_Perms(perms), eventhandler(0) {
	m_Buffer = new uint8_t[m_Size];
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
BackupInfo::~BackupInfo() {
	safeDeleteArray(m_Buffer);
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
bool BackupInfo::backup() {
	Debugger::debuggerBase->getState(m_State);
	return Debugger::debuggerBase->readBytes(m_Address, m_Buffer, m_Size);
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
bool BackupInfo::restore() {
	Debugger::debuggerBase->setState(m_State);
	return Debugger::debuggerBase->writeBytes(m_Address, m_Buffer, m_Size);
}

//------------------------------------------------------------------------------
// Name: MemoryRegions(edb::pid_t pid)
// Desc: constructor
//------------------------------------------------------------------------------
MemoryRegions::MemoryRegions() : QAbstractItemModel(0), m_PID(0) {
	
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
void MemoryRegions::setPID(edb::pid_t pid) {
	m_PID = pid;
	m_Regions.clear();
	sync();
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
void MemoryRegions::clear() {
	m_PID = 0;
	m_Regions.clear();
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
MemoryRegions::~MemoryRegions() {

}

//------------------------------------------------------------------------------
// Name: proccessMapLine(const QString &line, Region &region)
// Desc: parses the data from a line of a memory map file
//------------------------------------------------------------------------------
bool MemoryRegions::proccessMapLine(const QString &line, Region &region) {

	// TODO: use QT's way of doing this

	bool ret = false;
	
	std::istringstream ss(qPrintable(line));
	std::string perms;

	// read the starting address in hex
	if(ss >> std::hex >> region.start) {
	
		// skip the dash
		if(ss.ignore(1)) {
		
			// read the ending address and then the permissions
			if(ss >> std::hex >> region.end >> perms) {

				std::string temp[2];
				std::string name;
				
				// read in the region base, then skip the next two fields
				if(ss >> region.base >> temp[0] >> temp[1]) {
				
					// read in the name, we dont check for success because it's
					// allowed to not be there (and would therefore be the 
					// empty string)
					//ss >> name;
					
					char buf[256];
					ss.getline(buf, sizeof(buf), '\n');
					name = buf;

					// create a bitmask for the permissions
					region.m_Permissions = 0;
					if(perms[0] == 'r') region.m_Permissions |= Region::READ;
					if(perms[1] == 'w') region.m_Permissions |= Region::WRITE;
					if(perms[2] == 'x') region.m_Permissions |= Region::EXEC;

					region.name = QString(name.c_str()).trimmed();

					ret = true;
					
					// TODO: make this happen less often!
					if(!region.name.isEmpty()) {
						if(region.base == 0x00000000) {
							if(region.executable()) {
								Debugger::symbolManager().loadSymbolFile(region.name, region.start);
							}
						}
					}
				}
			}
		}
	}
	return ret;
}

//------------------------------------------------------------------------------
// Name: sync()
// Desc: reads a memory map file line by line
//------------------------------------------------------------------------------
void MemoryRegions::sync() {

	m_Regions.clear();

	if(m_PID != 0) {
#ifdef Q_OS_LINUX
		const QString mapFile(QString().sprintf("/proc/%d/maps", m_PID));
		
		QFile file(mapFile);
        if(file.open(QIODevice::ReadOnly | QIODevice::Text)) {
		
			QTextStream in(&file);
			QString line = in.readLine();
		
			while (!line.isNull()) {
				Region region;
				if(proccessMapLine(line, region)) {
					m_Regions.push_back(region);
				}
				line = in.readLine();
			}
		}
#else
	#error "Unsupported Platform"
#endif
	}
	reset();
}

//------------------------------------------------------------------------------
// Name: handleEvent(DebugEvent event)
// Desc: 
//------------------------------------------------------------------------------
void MemoryRegions::Region::handleEvent(DebugEvent event) {
	Q_UNUSED(event);
	
	// TODO: check that the event was caused by a hlt in our shellcode
	
	if(m_BackupInfo != 0) {
	
		// restore the original code and register state
		m_BackupInfo->restore();

		// update permissions mask
		m_Permissions = m_BackupInfo->perms();

		// restore the event handler
		Debugger::setDebugEventHandler(m_BackupInfo->eventhandler);

		safeDelete(m_BackupInfo);
	}
	
	m_Lock = 0;
}

//------------------------------------------------------------------------------
// Name: ~Region()
// Desc: 
//------------------------------------------------------------------------------
MemoryRegions::Region::~Region() {	
	safeDelete(m_BackupInfo);
}


//------------------------------------------------------------------------------
// Name: setPermissions(bool read, bool write, bool execute)
// Desc: attempts to change the permissions of a given region by injecting
//       shellcode into the debugged process, and executing a mprotect syscall
//       we probably need to ensure that the debug event server thread is not
//       waiting for an event at this moment because we need to catch the "hlt"
//       instruction we injected in order to restore the original code
//       we also probably need this code to work perfectly in order to implement
//       memory access breakpoints
//------------------------------------------------------------------------------
void MemoryRegions::Region::setPermissions(bool read, bool write, bool execute, edb::address_t tempAddress) {

#ifdef Q_OS_LINUX
	uint32_t perms = 0;
	const uint32_t len = size();
	const uint32_t addr = start;
	static const uint32_t syscallnum = __NR_mprotect;
	
	if(read)	perms |= READ;
	if(write)	perms |= WRITE;
	if(execute)	perms |= EXEC;
	
	// start nowhere near portable code
	#define LEN_OFFSET		1
	#define ADDR_OFFSET		6
	#define PERM_OFFSET		11
	#define SYSCALL_OFFSET	16
	
	// TODO: only compile this code on x86
	char shellcode[] = {
		"\xb9\x00\x10\x00\x00"	// mov    $0x00001000,%ecx
		"\xbb\x00\x40\x00\x00"	// mov    $0x00004000,%ebx
		"\xba\x07\x00\x00\x00"	// mov    $0x00000007,%edx
		"\xb8\x7d\x00\x00\x00"	// mov    $0x7d,%eax
		"\xcd\x80"				// int    $0x80
		"\xf4"					// hlt
	};
			
	// put the data in the right part of the shellcode
	std::memcpy(&shellcode[LEN_OFFSET], &len, sizeof(len));
	std::memcpy(&shellcode[ADDR_OFFSET], &addr, sizeof(addr));
	std::memcpy(&shellcode[PERM_OFFSET], &perms, sizeof(perms));
	std::memcpy(&shellcode[SYSCALL_OFFSET], &syscallnum, sizeof(syscallnum));
	
	#undef LEN_OFFSET
	#undef ADDR_OFFSET
	#undef PERM_OFFSET
	#undef SYSCALL_OFFSET
	// end nowhere near portable code
	
	m_BackupInfo = new BackupInfo(sizeof(shellcode), tempAddress, perms);

	if(m_BackupInfo->backup()) {
	
		// write out our shellcode
		if(Debugger::debuggerBase->writeBytes(tempAddress, shellcode, sizeof(shellcode))) {

			// set eip to our temp address
			State state;
			Debugger::debuggerBase->getState(state);
			state.eip = tempAddress;
			Debugger::debuggerBase->setState(state);

			m_Lock = 1;

			m_BackupInfo->eventhandler = Debugger::setDebugEventHandler(this);

			// run and wait for the "crash" caused by the hlt instruction
			// should be a SIGSEGV on Linux
			Debugger::debuggerBase->resume(0);

			// this is a crappy solution but it works good enough for now		
			while(m_Lock) {
				qApp->processEvents();
			}
			return;
		}
	}
	
	
	// we only get here if something was not successful, start cleaning up	
	safeDelete(m_BackupInfo);
#else
	#error "Unsupported Platform"
#endif
}

//------------------------------------------------------------------------------
// Name: setPermissions(bool read, bool write, bool execute)
// Desc: wrapper around normal set permissions
//------------------------------------------------------------------------------
void MemoryRegions::Region::setPermissions(bool read, bool write, bool execute) {
#ifdef Q_OS_LINUX
	edb::address_t tempAddress = 0x00000000;
	int count = 0;
	int ret = QMessageBox::Yes;
	const QVector<Region> regions = Debugger::memoryRegions().regions();
	
	// search for an executable region to run our shell code
	foreach(Region region, regions) {
		if(region.executable()) {
			if(tempAddress == 0x00000000) {
				tempAddress = region.start;
			}
			++count;
		}
	}
	
	if(count == 1 && !execute) {
		ret = QMessageBox::question(Debugger::debuggerGUI, "Removing Execute Permissions On Last Executable Region",
			"You are about to remove execute permissions from the last executable region. Because of the need "
			"to run code in the process to change permissions, there will be no way to undo this. In addition, "
			"the process will no longer be able to run as it will have no execute permissions in any regions. "
			"Odds are this is not what you want to do."
			"Are you sure you want to remove execute permissions from this region?",
			QMessageBox::Yes, QMessageBox::No, QMessageBox::NoButton);
	}
	

	if(ret == QMessageBox::Yes) {
		if(tempAddress != 0x00000000) {
			setPermissions(read, write, execute, tempAddress);
		} else {
			QMessageBox::information(Debugger::debuggerGUI, "No Suitable Address Found", "This feature relies on running shellcode in the debugged process, no executable memory region was found. Unfortunately, this means that no more region permission changes can be made (it also means that there is nothing the process can continue to do since it cannot execute at all).");
		}
	}
#else
	#error "Unsupported Platform"
#endif
}

//------------------------------------------------------------------------------
// Name: findBaseAddress(const QString &name, bool *ok) const
// Desc: 
//------------------------------------------------------------------------------
edb::address_t MemoryRegions::findBaseAddress(const QString &name, bool *ok) const {
	
	setBool(ok, false);
	
	foreach(Region region, m_Regions) {
		if(region.base == 0x00000000) {
			if(region.name.endsWith(name)) {
				setBool(ok, true);
				return region.start;
			}
		}
	}
	return 0x00000000;
}

//------------------------------------------------------------------------------
// Name: findRegion(edb::address_t address) const
// Desc: 
//------------------------------------------------------------------------------
bool MemoryRegions::findRegion(edb::address_t address) const {
	foreach(Region i, m_Regions) {
		if(address >= i.start && address < i.end) {
			return true;
		}
	}
	return false;
}

//------------------------------------------------------------------------------
// Name: findRegion(edb::address_t address, MemoryRegions::Region &region) const
// Desc: 
//------------------------------------------------------------------------------
bool MemoryRegions::findRegion(edb::address_t address, MemoryRegions::Region &region) const {

	foreach(Region i, m_Regions) {
		if(address >= i.start && address < i.end) {
			region = i;
			return true;
		}
	}
	return false;
}


//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
QVariant MemoryRegions::data(const QModelIndex &index, int role) const {
	if (!index.isValid())
		return QVariant();

	if (role != Qt::DisplayRole)
		return QVariant();

	switch(index.column()) {
	case 0:
		return QString().sprintf("%08x", m_Regions[index.row()].start);
	case 1:
		return QString().sprintf("%08x", m_Regions[index.row()].end);
	case 2:
		return QString().sprintf("%c%c%c", m_Regions[index.row()].readable() ? 'r' : '-', m_Regions[index.row()].writable() ? 'w' : '-', m_Regions[index.row()].executable() ? 'x' : '-' );
	case 3:
		return m_Regions[index.row()].name;
	default:
		return QVariant();
	}
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
QModelIndex MemoryRegions::index(int row, int column, const QModelIndex &parent) const {
	Q_UNUSED(parent);

	if(row >= m_Regions.size()) {
		return QModelIndex();
	}

	if(column >= 4) {
		return QModelIndex();
	}
	
	return createIndex(row, column, const_cast<Region *>(&m_Regions[row]));
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
QModelIndex MemoryRegions::parent(const QModelIndex &index) const {
	Q_UNUSED(index);
	return QModelIndex(); 
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
int MemoryRegions::rowCount(const QModelIndex &parent) const {
	Q_UNUSED(parent);
	return m_Regions.size();
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
int MemoryRegions::columnCount(const QModelIndex &parent) const {
	Q_UNUSED(parent);
	return 4;
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
QVariant MemoryRegions::headerData(int section, Qt::Orientation orientation, int role) const {
	if(role == Qt::DisplayRole && orientation == Qt::Horizontal) {
		switch(section) {
		case 0:
			return "Start Address";
		case 1:
			return "End Address";
		case 2:
			return "Permissions";
		case 3:
			return "Name";
		default:
			return QVariant();
		}	
	}
	
	return QVariant();
}

