/*
Copyright (C) 2006 - 2009 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 "MemRegion.h"
#include "Types.h"
#include "Util.h"
#include "State.h"
#include "DebuggerCoreInterface.h"
#include "DebugEventHandlerInterface.h"
#include "MemoryRegions.h"
#include "Debugger.h"

#include <QtGlobal>
#include <QMessageBox>
#include <QApplication>
#include <cstring>

#include <asm/unistd.h>
#include <sys/mman.h>

template <size_t N>
class BackupInfo : public DebugEventHandlerInterface {
public:
	BackupInfo(edb::address_t address, quint8 perms);
	
private:
	Q_DISABLE_COPY(BackupInfo)
	
public:
	quint8 perms() const	{ return m_Perms; }
	bool locked() const		{ return m_Lock != 0; }
	
public:
	bool backup();
	bool restore();
	
public:
	virtual void handleEvent(const DebugEvent &event);

private:
	volatile int			m_Lock;
	edb::address_t			m_Address;
	quint8					m_Perms;
	State					m_State;
	quint8					m_Buffer[N];
	
public:
	DebugEventHandlerInterface	*eventhandler;
};

//------------------------------------------------------------------------------
// Name: BackupInfo(td::size_t size, edb::address_t address, quint8 perms)
// Desc: 
//------------------------------------------------------------------------------
template <size_t N>
BackupInfo<N>::BackupInfo(edb::address_t address, quint8 perms) :
		m_Lock(1), m_Address(address), m_Perms(perms), eventhandler(0) {
}

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

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

//------------------------------------------------------------------------------
// Name: handleEvent(const DebugEvent &event)
// Desc: 
//------------------------------------------------------------------------------
template <size_t N>
void BackupInfo<N>::handleEvent(const DebugEvent &event) {
	Q_UNUSED(event);
	
	// TODO: check that the event was caused by a hlt in our shellcode	
	m_Lock = 0;
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
bool MemRegion::accessible() const {
	return readable() || writable() || executable();
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
bool MemRegion::readable() const {
	return (m_Permissions & PROT_READ) != 0;
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
bool MemRegion::writable() const {
	return (m_Permissions & PROT_WRITE) != 0;
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
bool MemRegion::executable() const {
	return (m_Permissions & PROT_EXEC) != 0;
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
int MemRegion::size() const {
	return end - start;
}

//------------------------------------------------------------------------------
// 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 MemRegion::setPermissions(bool read, bool write, bool execute, edb::address_t tempAddress) {

	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 |= PROT_READ;
	if(write)	perms |= PROT_WRITE;
	if(execute)	perms |= PROT_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
	
	typedef BackupInfo<sizeof(shellcode)> BI;
	try {
		BI backup_info(tempAddress, perms);

		if(backup_info.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);
				state.setInstructionPointer(tempAddress);
				edb::v1::debuggerBase->setState(state);

				backup_info.eventhandler = edb::v1::setDebugEventHandler(&backup_info);

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

				// TODO: this is a crappy solution but it works good enough for now		
				while(backup_info.locked()) {
					QCoreApplication::processEvents(QEventLoop::WaitForMoreEvents);
				}

				// restore the original code and register state
				backup_info.restore();

				// update permissions mask
				m_Permissions = backup_info.perms();

				// restore the event handler
				edb::v1::setDebugEventHandler(backup_info.eventhandler);
			}
		}
	} catch(const std::bad_alloc &) {
		QMessageBox::information(0, QObject::tr("Memroy Allocation Error"),
			QObject::tr("Unable to satisfy memory allocation request for backup code."));
	}	
}

//------------------------------------------------------------------------------
// Name: ~MemRegion()
// Desc: 
//------------------------------------------------------------------------------
MemRegion::~MemRegion() {
}

//------------------------------------------------------------------------------
// Name: setPermissions(bool read, bool write, bool execute)
// Desc: wrapper around normal set permissions
//------------------------------------------------------------------------------
void MemRegion::setPermissions(bool read, bool write, bool execute) {
	edb::address_t tempAddress			= edb::nullptr;
	int count							= 0;
	int ret								= QMessageBox::Yes;
	const MemoryRegions &memoryRegions	= edb::v1::memoryRegions();
	const QVector<MemRegion> &regions	= memoryRegions.regions();
	
	// search for an executable region to run our shell code
	Q_FOREACH(const MemRegion &region, regions) {
		if(region.executable()) {
			if(tempAddress == edb::nullptr) {
				tempAddress = region.start;
			}
			
			if(++count > 1) {
				break;
			}
		}
	}
	
	if(executable() && count == 1 && !execute) {
		ret = QMessageBox::question(0, 
			QObject::tr("Removing Execute Permissions On Last Executable MemRegion"),
			QObject::tr("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);
	}
	
	if(ret == QMessageBox::Yes) {
		if(tempAddress != edb::nullptr) {
			setPermissions(read, write, execute, tempAddress);
		} else {
			QMessageBox::information(0,
			QObject::tr("No Suitable Address Found"),
			QObject::tr("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)."));
		}
	}
}
