/*
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 <cstring>
#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, quint8 perms);
	~BackupInfo();
	
private:
	BackupInfo &operator=(const BackupInfo &);
	BackupInfo(const BackupInfo&);
	
public:
	uint perms() const { return m_Perms; }
	
public:
	bool backup();
	bool restore();
	
private:
	quint8			*m_Buffer;
	std::size_t		m_Size;
	edb::address_t	m_Address;
	uint			m_Perms;
	State			m_State;
	
public:
	DebugEventHandlerInterface	*eventhandler;
};

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

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

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

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

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

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

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

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

//------------------------------------------------------------------------------
// 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) {

	bool ret = false;
#if defined (Q_OS_LINUX)
	const QStringList items = line.split(" ", QString::SkipEmptyParts);
	if(items.size() >= 3) {
		bool ok;
		const QStringList bounds = items[0].split("-");
		if(bounds.size() == 2) {
			region.start = bounds[0].toULongLong(&ok, 16);
			if(ok) {
				region.end = bounds[1].toULongLong(&ok, 16);
				if(ok) {
					region.base = items[2].toULongLong(&ok, 16);
					if(ok) {
						const QString perms	= items[1];
						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;

						if(items.size() >= 6) {
							region.name				= items[5];
						}

						// if the region has a name, is mapped starting
						// at the begining of the file, and is executable, sounds 
						// like a module mapping!
						if(!region.name.isEmpty()) {
							if(region.base == 0x00000000) {
								if(region.executable()) {
									edb::v1::symbolManager().loadSymbolFile(region.name, region.start);
								}
							}
						}

						ret = true;
					}
				}
			}
		}
	}
#elif defined (Q_OS_FREEBSD)

	// TODO: use QT's way of doing this
	std::istringstream ss(qPrintable(line));
	std::string perms;
	
	std::string temp[9]
	//0x8048000 0x8057000 13 0 0xc27e64a4 r-x 2 0 0x0 COW NC vnode /usr/local/libexec/gam_server

	// read the starting address in hex
	if(ss >> std::hex >> region.start) {
		
		// read the ending address and then the permissions
		if(ss >> std::hex >> region.end) {
			if(ss >> temp[0] >> temp[1] >> temp[2] >> perms) {
				if(ss >> temp[3] >> temp[4] >> temp[5] >> temp[6] >> temp[7]>> temp[8]) {
					std::string 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;
				}
			}
		}
	}
#else
	#error "Unsupported Platform"
#endif
	return ret;
}

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

	m_Regions.clear();

	if(m_PID != 0) {
#if defined (Q_OS_LINUX)
		const QString mapFile(QString().sprintf("/proc/%d/maps", m_PID));
#elif defined (Q_OS_FREEBSD)
		const QString mapFile(QString().sprintf("/proc/%d/map", m_PID));
#else
	#error "Unsupported Platform"
#endif
	
		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();
			}
		}
	}
	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
		edb::v1::setDebugEventHandler(m_BackupInfo->eventhandler);

		release(m_BackupInfo);
	}
	
	m_Lock = 0;
}

//------------------------------------------------------------------------------
// Name: ~Region()
// Desc: 
//------------------------------------------------------------------------------
MemoryRegions::Region::~Region() {	
	release(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

	edb::address_t perms = 0;
	const edb::address_t len = size();
	const edb::address_t addr = start;
	static const edb::address_t syscallnum = __NR_mprotect;
	
	if(read)	perms |= READ;
	if(write)	perms |= WRITE;
	if(execute)	perms |= EXEC;
	
	// start of nowhere near portable code
#if defined(EDB_X86)
	#define ADDR_OFFSET		1
	#define LEN_OFFSET		6
	#define PERM_OFFSET		11
	#define SYSCALL_OFFSET	16
	
	quint8 shellcode[] = {
		"\xbb\x00\x40\x00\x00"	// mov    $0x00004000,%ebx
		"\xb9\x00\x10\x00\x00"	// mov    $0x00001000,%ecx
		"\xba\x07\x00\x00\x00"	// mov    $0x00000007,%edx
		"\xb8\x7d\x00\x00\x00"	// mov    $0x7d,%eax
		"\xcd\x80"				// int    $0x80
		"\xf4"					// hlt
	};
#elif defined(EDB_X86_64)
	#define ADDR_OFFSET		2
	#define LEN_OFFSET		12
	#define PERM_OFFSET		22
	#define SYSCALL_OFFSET	32

	quint8 shellcode[] = {
		"\x48\xbf\x22\x22\x22\x22\x11\x11\x11\x11"	// mov    $0x1111111122222222,%rdi
		"\x48\xbe\x44\x44\x44\x44\x33\x33\x33\x33"	// mov    $0x3333333344444444,%rsi
		"\x48\xba\x66\x66\x66\x66\x55\x55\x55\x55"	// mov    $0x5555555566666666,%rdx
		"\x48\xb8\x88\x88\x88\x88\x77\x77\x77\x77"	// mov    $0x7777777788888888,%rax
		"\x0f\x05"									// syscall
		"\xf4"										// hlt
	};
#else
#error "invalid architecture"
#endif
	// 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(edb::v1::debuggerBase->writeBytes(tempAddress, shellcode, sizeof(shellcode))) {

			// set eip to our temp address
			State state;
			edb::v1::debuggerBase->getState(state);
#if defined(EDB_X86)
			state.eip = tempAddress;
#elif defined(EDB_X86_64)
			state.rip = tempAddress;
#else
#error "invalid architecture"
#endif
			edb::v1::debuggerBase->setState(state);

			m_Lock = 1;

			m_BackupInfo->eventhandler = edb::v1::setDebugEventHandler(this);

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

			// TODO: this is a crappy solution but it works good enough for now		
			while(m_Lock) {
				QCoreApplication::processEvents(QEventLoop::WaitForMoreEvents);
			}
			return;
		}
	}
	
	
	// we only get here if something was not successful, start cleaning up	
	release(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 = edb::v1::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(edb::v1::debuggerUI, "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(edb::v1::debuggerUI, "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: 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() && role == Qt::DisplayRole) {
	
		const MemoryRegions::Region &region = m_Regions[index.row()];
		QString temp;
	
		switch(index.column()) {
		case 0:
			return temp.sprintf(EDB_FMT_PTR, region.start);
		case 1:
			return temp.sprintf(EDB_FMT_PTR, region.end);
		case 2:
			return QString("%1%2%3").arg(region.readable() ? 'r' : '-').arg(region.writable() ? 'w' : '-').arg(region.executable() ? 'x' : '-' );
		case 3:
			return region.name;
		}
	}
	
	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";
		}	
	}
	
	return QVariant();
}

