/*
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 "DebuggerCoreWindows.h"
#include "DebugEvent.h"
#include "State.h"
#include <QDebug>
#include <QStringList>
#include <windows.h>


class Win32Handle {
public:
	Win32Handle() : handle(NULL) {}
	Win32Handle(HANDLE handle) : handle(handle) {}
	~Win32Handle() { if(handle) { CloseHandle(handle); } }
private:
	Win32Handle(const Win32Handle&);
	Win32Handle &operator=(const Win32Handle&);

public:
	operator HANDLE() const { return handle; }

private:
	HANDLE handle;
};

//------------------------------------------------------------------------------
// Name: DebuggerCore()
// Desc: constructor
//------------------------------------------------------------------------------
DebuggerCore::DebuggerCore() : m_PageSize(0), m_ProcessHandle(0){
	DebugSetProcessKillOnExit(false);
	
	SYSTEM_INFO sys_info;
	GetSystemInfo(&sys_info);
	m_PageSize = sys_info.dwPageSize;
}

//------------------------------------------------------------------------------
// Name: 
// Desc: 
//------------------------------------------------------------------------------
DebuggerCore::~DebuggerCore() {
	detach();
}

//------------------------------------------------------------------------------
// Name: pageSize() const
// Desc: returns the size of a page on this system
//------------------------------------------------------------------------------
edb::address_t DebuggerCore::pageSize() const {
	return m_PageSize;
}

//------------------------------------------------------------------------------
// Name: waitForDebugEvent(DebugEvent &event, bool &ok, int secs)
// Desc: waits for a debug event, secs is a timeout (but is not yet respected)
//       ok will be set to false if the timeout expires
//------------------------------------------------------------------------------
bool DebuggerCore::waitForDebugEvent(DebugEvent &event, int msecs) {
	if(attached()) {
		DEBUG_EVENT de;
		while(WaitForDebugEvent(&de, msecs == 0 ? INFINITE : msecs) != 0) {
			switch(de.dwDebugEventCode) {
			case CREATE_THREAD_DEBUG_EVENT:
				m_Threads.push_back(de.dwThreadId);
				ContinueDebugEvent(de.dwProcessId, de.dwThreadId, DBG_CONTINUE);
				break;
			case EXIT_THREAD_DEBUG_EVENT:
				m_Threads.removeAll(de.dwThreadId);
				ContinueDebugEvent(de.dwProcessId, de.dwThreadId, DBG_CONTINUE);
				break;
			case CREATE_PROCESS_DEBUG_EVENT:
			case EXIT_PROCESS_DEBUG_EVENT:
			case LOAD_DLL_DEBUG_EVENT:
			case UNLOAD_DLL_DEBUG_EVENT:
			case OUTPUT_DEBUG_STRING_EVENT:
			case RIP_EVENT:
				ContinueDebugEvent(de.dwProcessId, de.dwThreadId, DBG_CONTINUE);
				break;
			case EXCEPTION_DEBUG_EVENT:
				event = DebugEvent(de);
				m_ActiveThread = event.thread();
				return true;
			}
		}
	}
	return false;
}

//------------------------------------------------------------------------------
// Name: readPages(edb::address_t address, void *buf, std::size_t count)
// Desc: reads <count> pages from the process starting at <address>
// Note: buf's size must be >= count * pageSize()
// Note: address should be page aligned.
//------------------------------------------------------------------------------
bool DebuggerCore::readPages(edb::address_t address, void *buf, std::size_t count) {

	Q_ASSERT(buf != 0);
	
	bool ok = false;

	if(attached()) {
		const edb::address_t orig_address	= address;
		const edb::address_t page_size		= pageSize();
		const edb::address_t end_address	= orig_address + page_size * count;

		if(m_ProcessHandle != 0) {
			SIZE_T bytes_read;
			ok = ReadProcessMemory(
				m_ProcessHandle,
				reinterpret_cast<LPVOID>(address),
				buf,
				count * pageSize(),
				&bytes_read);
		}
		
		if(ok) {
			quint8 *const orig_ptr = reinterpret_cast<quint8 *>(buf);
			// TODO: handle if breakponts have a size more than 1!
			Q_FOREACH(const QSharedPointer<Breakpoint> &bp, m_Breakpoints) {		
				if(bp->address() >= orig_address && bp->address() < end_address) {
					// show the original bytes in the buffer..
					orig_ptr[bp->address() - orig_address] = bp->originalBytes()[0];
				}
			}
		}
	}
	return ok;
}

//------------------------------------------------------------------------------
// Name: readBytes(edb::address_t address, void *buf, std::size_t len)
// Desc: reads <len> bytes into <buf> starting at <address>
// Note: if the read failed, the part of the buffer that could not be read will 
//       be filled with 0xff bytes
//------------------------------------------------------------------------------
bool DebuggerCore::readBytes(edb::address_t address, void *buf, std::size_t len) {

	Q_ASSERT(buf != 0);

	bool ok = false;
	
	if(attached()) {
		const edb::address_t orig_address	= address;
		const edb::address_t end_address	= orig_address + len;

		if(m_ProcessHandle != 0) {
			SIZE_T bytes_read;
			ok = ReadProcessMemory(
				m_ProcessHandle,
				reinterpret_cast<LPVOID>(address),
				buf,
				len,
				&bytes_read);
		}
		// TODO: do the documented 0xff fill
		
		if(ok) {
			quint8 *const orig_ptr = reinterpret_cast<quint8 *>(buf);
			// TODO: handle if breakponts have a size more than 1!
			Q_FOREACH(const QSharedPointer<Breakpoint> &bp, m_Breakpoints) {		
				if(bp->address() >= orig_address && bp->address() < end_address) {
					// show the original bytes in the buffer..
					orig_ptr[bp->address() - orig_address] = bp->originalBytes()[0];
				}
			}
		}
	}
	return ok;
}

//------------------------------------------------------------------------------
// Name: writeBytes(edb::address_t address, const void *buf, std::size_t len)
// Desc: writes <len> bytes from <buf> starting at <address>
//------------------------------------------------------------------------------
bool DebuggerCore::writeBytes(edb::address_t address, const void *buf, std::size_t len) {
	
	Q_ASSERT(buf != 0);
	
	bool ok = false;
	if(attached()) {
		if(m_ProcessHandle != 0) {
			SIZE_T bytes_written;
			ok = WriteProcessMemory(
				m_ProcessHandle,
				reinterpret_cast<LPVOID>(address),
				buf,
				len,
				&bytes_written);
			
			if(ok) {
				FlushInstructionCache(m_ProcessHandle, reinterpret_cast<LPVOID>(address), bytes_written);
			}
		}
	}
	return ok;
}

//------------------------------------------------------------------------------
// Name: attach(edb::pid_t pid)
// Desc: 
//------------------------------------------------------------------------------
bool DebuggerCore::attach(edb::pid_t pid) {
	detach();
	
	if(DebugActiveProcess(pid)) {
		DEBUG_EVENT ev;
		if(WaitForDebugEvent(&ev, INFINITE)) {

			assert(ev.dwDebugEventCode == CREATE_PROCESS_DEBUG_EVENT);

			// close the image file handle
			CloseHandle(ev.u.CreateProcessInfo.hFile);

			// remember the process handle
			m_ProcessHandle = ev.u.CreateProcessInfo.hProcess;

			// for now we'l close the thread's handle
			CloseHandle(ev.u.CreateProcessInfo.hThread);

			m_PID = pid;
			m_ActiveThread = ev.dwThreadId;

			// tell it to continue to we can catch the rest of the events
			ContinueDebugEvent(ev.dwProcessId, ev.dwThreadId, DBG_CONTINUE);

			return true;
		}
	}
	
	return false;
}

//------------------------------------------------------------------------------
// Name: detach()
// Desc: 
//------------------------------------------------------------------------------
void DebuggerCore::detach() {
	if(attached()) {
		clearBreakpoints();
		m_Threads.clear();

		// why is this illegal?
		//CloseHandle(m_ProcessHandle);
		m_ProcessHandle = 0;

		DebugActiveProcessStop(pid());
		m_PID = 0;
	}
}

//------------------------------------------------------------------------------
// Name: kill()
// Desc: 
//------------------------------------------------------------------------------
void DebuggerCore::kill() {
	if(attached()) {
		clearBreakpoints();
		if(m_ProcessHandle != 0) {
			TerminateProcess(m_ProcessHandle, -1);
			CloseHandle(m_ProcessHandle);
			m_ProcessHandle = 0;
		}
		m_PID = 0;
	}
}

//------------------------------------------------------------------------------
// Name: pause()
// Desc: stops *all* threads of a process
//------------------------------------------------------------------------------
void DebuggerCore::pause() {

	if(attached()) {
		if(m_ProcessHandle != 0) {
			DebugBreakProcess(m_ProcessHandle);
		}
	}
}

//------------------------------------------------------------------------------
// Name: resume(EVENT_STATUS status)
// Desc: 
//------------------------------------------------------------------------------
void DebuggerCore::resume(EVENT_STATUS status) {
	// TOOD: assert that we are paused
	
	if(attached()) {
		// TODO: does this resume *all* threads?
		ContinueDebugEvent(
			pid(),
			activeThread(),
			(status == DebuggerCore::DEBUG_CONTINUE) ? (DBG_CONTINUE) : (DBG_EXCEPTION_NOT_HANDLED));
	}
}

//------------------------------------------------------------------------------
// Name: step(EVENT_STATUS status)
// Desc: 
//------------------------------------------------------------------------------
void DebuggerCore::step(EVENT_STATUS status) {
	// TOOD: assert that we are paused
	
	if(attached()) {
		const Win32Handle th = OpenThread(THREAD_ALL_ACCESS, FALSE, activeThread());
		if(th != 0) {		
			CONTEXT	context;
			context.ContextFlags = CONTEXT_CONTROL;
			GetThreadContext(th, &context);
			
			context.EFlags |= (1 << 8); // set the TF bit
			SetThreadContext(th, &context);
			
			ContinueDebugEvent(
				pid(),
				activeThread(),
				(status == DebuggerCore::DEBUG_CONTINUE) ? (DBG_CONTINUE) : (DBG_EXCEPTION_NOT_HANDLED));
		}
	}
}

//------------------------------------------------------------------------------
// Name: getState(State &state)
// Desc: 
//------------------------------------------------------------------------------
void DebuggerCore::getState(State &state) {
	// TOOD: assert that we are paused
	
	if(attached()) {
		const Win32Handle th = OpenThread(THREAD_ALL_ACCESS, FALSE, activeThread());
		if(th != 0) {	
			CONTEXT	context;
			context.ContextFlags = CONTEXT_FULL | CONTEXT_DEBUG_REGISTERS | CONTEXT_FLOATING_POINT;
			GetThreadContext(th, &context);
			
			state.eax = context.Eax;
			state.ebx = context.Ebx;
			state.ecx = context.Ecx;
			state.edx = context.Edx;
			state.esp = context.Esp;
			state.ebp = context.Ebp;
			state.edi = context.Edi;
			state.esi = context.Esi;
			state.eip = context.Eip;
			state.eflags = context.EFlags;
			state._cs = context.SegCs;
			state._ds = context.SegDs;
			state._es = context.SegEs;
			state._fs = context.SegFs;
			state._gs = context.SegGs;
			state._ss = context.SegSs;
			state.orig_eax = 0;
			
			// TODO: FPU stuff
			
			state.dr[0] = context.Dr0;
			state.dr[1] = context.Dr1;
			state.dr[2] = context.Dr2;
			state.dr[3] = context.Dr3;
			state.dr[4] = 0;
			state.dr[5] = 0;
			state.dr[6] = context.Dr6;
			state.dr[7] = context.Dr7;

			LDT_ENTRY ldt_entry;
			if(GetThreadSelectorEntry(th, context.SegGs, &ldt_entry)) {
				state.gs_base = ldt_entry.BaseLow | (ldt_entry.HighWord.Bits.BaseMid << 16) | (ldt_entry.HighWord.Bits.BaseHi << 24);;
			}

			if(GetThreadSelectorEntry(th, context.SegFs, &ldt_entry)) {
				state.fs_base = ldt_entry.BaseLow | (ldt_entry.HighWord.Bits.BaseMid << 16) | (ldt_entry.HighWord.Bits.BaseHi << 24);;
			}	
		}
	} else {
		state.clear();
	}
}

//------------------------------------------------------------------------------
// Name: setState(const State &state)
// Desc: 
//------------------------------------------------------------------------------
void DebuggerCore::setState(const State &state) {
	
	// TOOD: assert that we are paused
	
	if(attached()) {
		CONTEXT	context;
		context.ContextFlags = CONTEXT_FULL | CONTEXT_DEBUG_REGISTERS | CONTEXT_FLOATING_POINT;
		
		context.Eax = state.eax;
		context.Ebx = state.ebx;
		context.Ecx = state.ecx;
		context.Edx = state.edx;
		context.Esp = state.esp;
		context.Ebp = state.ebp;
		context.Edi = state.edi;
		context.Esi = state.esi;
		context.Eip = state.eip;
		context.EFlags = state.eflags;
		context.SegCs = state._cs;
		context.SegDs = state._ds;
		context.SegEs = state._es;
		context.SegFs = state._fs;
		context.SegGs = state._gs;
		context.SegSs = state._ss;
		
		// TODO: FPU stuff
		
		context.Dr0 = state.dr[0];
		context.Dr1 = state.dr[1];
		context.Dr2 = state.dr[2];
		context.Dr3 = state.dr[3];
		//context.Dr4 = state.dr[4];
		//context.Dr5 = state.dr[5];
		context.Dr6 = state.dr[6];
		context.Dr7 = state.dr[7];

		const Win32Handle th = OpenThread(THREAD_ALL_ACCESS, FALSE, activeThread());
		if(th != 0) {
			SetThreadContext(th, &context);
		}
	}
}

//------------------------------------------------------------------------------
// Name: open(const QString &path, const QString &cwd, const QStringList &args, const QString &tty)
// Desc: 
//------------------------------------------------------------------------------
bool DebuggerCore::open(const QString &path, const QString &cwd, const QStringList &args, const QString &tty) {
	
	Q_UNUSED(path);
	Q_UNUSED(cwd);
	Q_UNUSED(args);
	Q_UNUSED(tty);

	detach();
	bool ret = false;

#if 0	
	STARTUPINFO			startup_info;
	PROCESS_INFORMATION	process_info;
	
	std::memset(&startup_info, 0, sizeof(startup_info));
	std::memset(&process_info, 0, sizeof(process_info));
	
	if(CreateProcess(
			0,		// exe
			0,		// commandline
			NULL,	// default security attributes
			NULL,	// default thread security too
			FALSE,	// inherit handles
			CREATE_SUSPENDED | DEBUG_ONLY_THIS_PROCESS,
			NULL,	// environment
			0,		// cwd
			&startup_info,
			&process_info)) {

		m_PID			= process_info.dwProcessId;
		m_ActiveThread	= process_info.dwThreadId;

		CloseHandle(process_info.hProcess);
		CloseHandle(process_info.hThread);
		ret = true;
	}
#endif
	return ret;
}

Q_EXPORT_PLUGIN2(DebuggerCore, DebuggerCore)
